"""Dilated-convolution forecasting backbone. An attention-free encoder: a stack of causal dilated convolutions. With kernel K=3 and dilations d_i = 2^(i-1) for i=1..N the receptive field is RF = 1 + 2 * sum_i d_i = 1 + 2 * (2^N - 1) so N=10 layers cover RF=2047, sufficient for the L=2048 context with zero downsampling and no information loss at any time scale. Native multi-scale via the dilation schedule; deployment-friendly (no L^2 attention matrix, no softmax, pure matmul + element-wise; quantizes cleanly to INT8); streaming- friendly (left-only causal padding). Structural priors (zero-parameter): - a normalized-periodogram period detector driving a phase encoding - bounded recency basis (signed-linear/log, multi-scale exp decay) - position-parameterized decoder queries (single-shot arbitrary horizon) - no autoregressive rollout inside the backbone """ from __future__ import annotations import math import torch import torch.nn as nn import torch.nn.functional as F from .periodogram import significant_periods from .encoding import ( N_RECENCY_CHANNELS, _norm_fp32, _phase_encoding, _positional_encoding, ) class _SwiGLU(nn.Module): """Standard SwiGLU FFN.""" def __init__(self, d: int, d_hidden: int) -> None: super().__init__() self.up = nn.Linear(d, 2 * d_hidden) self.down = nn.Linear(d_hidden, d) def forward(self, x: torch.Tensor) -> torch.Tensor: gate, val = self.up(x).chunk(2, dim=-1) return self.down(F.silu(gate) * val) class _DilatedConvBlock(nn.Module): """Dilated Conv1d → RMSNorm → SwiGLU → RMSNorm with residuals. ``causal=True``, which is what the released config sets, pads all (K-1)*d timesteps on the left, so each output position sees only its own past. With ``causal=False`` the padding is centered instead: length is still preserved and each position gets a symmetric view of (K-1)*d/2 timesteps on either side, at the cost of future-side context. Under both, the dilation scales the per-layer receptive field without adding parameters. """ def __init__( self, d: int, kernel: int = 3, dilation: int = 1, ffn_mult: float = 1.5, causal: bool = False, gated: bool = False, separable: bool = False, ) -> None: super().__init__() self.k = int(kernel) self.dilation = int(dilation) self.causal = bool(causal) self.gated = bool(gated) self.separable = bool(separable) # Conv with dilation; padding handled in forward. Centered padding # gives each position a symmetric view but feeds the right edge # FUTURE-side zeros: at the deepest layer the last (most recent) # position's representation is dominated by padding standing in for # the unknown forecast, and the model learns a train/inference # mismatch. Causal (all-left) padding removes both pathologies and # is a prerequisite for honest streaming inference. if self.separable: # Depthwise-separable factorization: depthwise (per-channel, dilated, # padding consumed in forward) + pointwise 1x1 (channel mix). Params # D*K + D*D against a full conv's D*D*K: the receptive field is # preserved at a fraction of the parameters per block. Padding # before self.conv feeds the depthwise stage. self.conv = nn.Sequential( nn.Conv1d(d, d, kernel_size=self.k, dilation=self.dilation, groups=d), nn.Conv1d(d, d, kernel_size=1), ) else: self.conv = nn.Conv1d(d, d, kernel_size=self.k, dilation=self.dilation) # Lightweight gated conv (multiplicative gating): a cheap # DEPTHWISE gate conv produces a sigmoid mask over the main conv output, # conv_out * σ(gate). Adds data-dependent gating to the encoder at ~D·k # params/block (vs doubling the full conv). Off by default. self.gate = ( nn.Conv1d(d, d, kernel_size=self.k, dilation=self.dilation, groups=d) if self.gated else None ) self.norm1 = nn.RMSNorm(d) d_hidden = int(d * ffn_mult) self.ffn = _SwiGLU(d, d_hidden) self.norm2 = nn.RMSNorm(d) def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (B, L, D). Conv1d takes (B, D, L). x_t = x.transpose(1, 2) # Effective kernel span = (K-1)*d + 1. pad_total = (self.k - 1) * self.dilation if self.causal: left, right = pad_total, 0 else: left = pad_total // 2 right = pad_total - left x_t = F.pad(x_t, (left, right)) if self.separable and self.gate is None: # Depthwise (B,D,L); then the pointwise 1x1 conv == a per-position # Linear, run as F.linear (cublas GEMM) instead of an im2col/cudnn # 1x1-conv kernel. Bit-identical channel contraction; params stay as # conv.0/conv.1 so existing checkpoints load unchanged. dw = self.conv[0](x_t).transpose(1, 2) # (B, L, D) pw = self.conv[1] conv_out = F.linear(dw, pw.weight.squeeze(-1), pw.bias) # (B, L, D) else: conv_out = self.conv(x_t) if self.gate is not None: conv_out = conv_out * torch.sigmoid(self.gate(x_t)) conv_out = conv_out.transpose(1, 2) x = _norm_fp32(self.norm1, x + conv_out) x = _norm_fp32(self.norm2, x + self.ffn(x)) return x class DilatedConvBackbone(nn.Module): """Dilated-conv encoder with a phase-conditioned, position-parameterized decoder. Args: seq_len: L, the context length. p_out: H, the single-shot output length. n_quantiles: Q, the number of output channels. d: channel dimension. n_layers: number of dilated-conv blocks (10 → RF 2047 ≥ L=2048). kernel: conv kernel size (default 3). ffn_mult: SwiGLU hidden multiplier (default 1.5). dilations: explicit dilation schedule; if None, uses 2^(i-1). top_k_periods: the periodogram detector's top-K (default 4). significance_alpha: the periodogram detector's Bonferroni alpha (default 0.05). n_harmonics: Fourier harmonics per detected period (default 1). pool_kind: context-summary pooling: "mean_last" (default), "mean", "last". Concat the chosen pool(s) into the per-horizon query before query_proj. causal: if True, all conv padding is left-only (no future leakage, clean right edge, streaming-honest). phase_bins: if > 0, augment the global pool with a period-folded seasonal profile: for each period the periodogram detector returns, fold the encoder output into this many phase bins and let each decoder query gather the bin matching its own phase. 0 disables phase folding and leaves the plain global pool, which is the control the paper's phase-folding ablation is measured against. """ def __init__( self, seq_len: int, p_out: int, n_quantiles: int = 1, d: int = 80, n_layers: int = 10, kernel: int = 3, ffn_mult: float = 1.5, dilations: list[int] | None = None, top_k_periods: int = 4, significance_alpha: float = 0.05, n_harmonics: int = 1, pool_kind: str = "mean_last", causal: bool = False, phase_bins: int = 0, phase_stats: str = "mean", phase_recency_tau: float = 0.0, recency_bins: int = 0, sig_gate: bool = False, cross_cycle: bool = False, decoder_depth: int = 1, horizon_kernel: int = 0, horizon_recurrence: bool = False, min_cycles: int = 0, period_trust: str = "off", gated_conv: bool = False, residual_naive: bool = False, residual_multi: bool = False, residual_trend: bool = False, decompose_kernel: int = 0, periodogram_off: bool = False, res_adaptive: bool = False, res_period_target: int = 64, res_r_max: float = 32.0, with_missing: bool = False, missing_channel: bool = False, separable_conv: bool = False, share_ffn: bool = False, future_conv: bool = False, future_conv_layers: int = 6, future_conv_seed: int = 128, base_seasonality: float = 24.0, local_anchor: bool = False, ) -> None: super().__init__() self.L = int(seq_len) self.p_out = int(p_out) self.n_quantiles = int(n_quantiles) self.D = int(d) self.K = int(top_k_periods) self.significance_alpha = float(significance_alpha) self.n_harmonics = int(n_harmonics) self.pool_kind = str(pool_kind) self.causal = bool(causal) self.phase_bins = int(phase_bins) if phase_stats not in ("mean", "mean_var"): raise ValueError(f"phase_stats={phase_stats!r}; expected 'mean'|'mean_var'.") self.phase_stats = str(phase_stats) self.stat_mult = 2 if self.phase_stats == "mean_var" else 1 self.phase_recency_tau = float(phase_recency_tau) self.recency_bins = int(recency_bins) self.sig_gate = bool(sig_gate) self.cross_cycle = bool(cross_cycle) self.decoder_depth = max(1, int(decoder_depth)) self.horizon_kernel = int(horizon_kernel) self.horizon_recurrence = bool(horizon_recurrence) self.min_cycles = int(min_cycles) if period_trust not in ("off", "coverage", "full"): raise ValueError(f"period_trust={period_trust!r}; expected off|coverage|full") self.period_trust = str(period_trust) # Per-period reliability weight w_k = sigmoid(linear([margin, ln # coverage])): a continuous down-weighting of long or weakly-supported # periods, whose crossover is learned rather than set. It covers the # same ground as the integer min_cycles cutoff, which stays available # and independent; both are off in the released config. "coverage" uses # ln(L/period) alone (min_cycles is its hard-threshold limit); "full" # adds the significance margin ln(s_k/t_alpha) from the detector's # periodogram scores. if self.period_trust != "off": n_feat = 1 if self.period_trust == "coverage" else 2 self.pt = nn.Linear(n_feat, 1) with torch.no_grad(): self.pt.weight.zero_(); self.pt.weight[0, 0] = 1.0 # cov (or margin) coeff = 1 self.pt.bias.zero_() # data-determined Bonferroni threshold t_alpha (no tuned knob) n_fft = 1 << int(math.ceil(math.log2(max(2, self.L)))) self._n_bins = max(2, n_fft // 2) self._t_alpha = math.log(self._n_bins / max(self.significance_alpha, 1e-12)) / self._n_bins else: self.pt = None self.gated_conv = bool(gated_conv) self.residual_naive = bool(residual_naive) self.residual_multi = bool(residual_multi) self.residual_trend = bool(residual_trend) self.periodogram_off = bool(periodogram_off) self.res_adaptive = bool(res_adaptive) self.res_period_target = int(res_period_target) self.res_r_max = float(res_r_max) # Series decomposition: moving-avg trend / seasonal split fed # as two input channels. 0 disables. Even kernel → +1 for centered. self.decompose_kernel = int(decompose_kernel) self.with_missing = bool(with_missing) # Missing-value channel: feed the encoder a binary observed-mask so it # can distinguish a genuinely-unobserved position from a real value (the # faithful treatment, vs mean-fill which conflates the two). self.missing_channel = bool(missing_channel) if self.missing_channel and bool(res_adaptive): raise NotImplementedError( "missing_channel + res_adaptive: the observed-mask is not warped " "through _resolution_adapt; not supported together." ) # Recent-anchor channels: gradient-connected causal local level/scale so the # encoder can re-anchor amplitude under non-stationarity (the denorm origin # x_min is a frozen detached GLOBAL min, with no learned path to the # recent regime). self.local_anchor = bool(local_anchor) if self.n_harmonics < 1: raise ValueError(f"n_harmonics must be >= 1; got {self.n_harmonics}") if pool_kind not in ("mean_last", "mean", "last"): raise ValueError( f"pool_kind={pool_kind!r}; expected 'mean_last' | 'mean' | 'last'." ) # Positional encoding: phase channels + bounded recency channels. n_phase = 2 * self.K * self.n_harmonics n_pe = n_phase + N_RECENCY_CHANNELS # Input channels: raw value, or [trend, seasonal] if decomposing. n_value_ch = 2 if self.decompose_kernel > 0 else 1 if self.missing_channel: n_value_ch += 1 # +observed-mask if self.local_anchor: n_value_ch += 2 # +[local-scale residual, log local-scale] in_channels = n_value_ch + n_pe self.in_proj = nn.Linear(in_channels, self.D) if self.local_anchor: # zero-init the 2 anchor columns (last of the value channels) so the # model is baseline-equivalent at init and learns the anchor from zero. with torch.no_grad(): self.in_proj.weight[:, n_value_ch - 2:n_value_ch].zero_() # Dilation schedule. if dilations is None: dilations = [2**i for i in range(int(n_layers))] if len(dilations) != int(n_layers): raise ValueError( f"dilations length {len(dilations)} != n_layers {n_layers}" ) self.dilations = list(dilations) # Receptive field sanity (informational only; fails soft if RF < L). rf = 1 + (kernel - 1) * sum(self.dilations) self.receptive_field = rf self.encoder = nn.ModuleList([ _DilatedConvBlock( self.D, kernel=int(kernel), dilation=int(d_i), ffn_mult=float(ffn_mult), causal=self.causal, gated=self.gated_conv, separable=bool(separable_conv), ) for d_i in self.dilations ]) # Cross-layer FFN weight sharing (weight-tied): the SwiGLU FFN is the # largest param bucket and is dilation-independent, so one shared FFN # across all blocks recovers ~(n_layers-1)/n_layers of FFN params. The # per-block dilated convs (which carry the receptive field) stay distinct. self.share_ffn = bool(share_ffn) if self.share_ffn and len(self.encoder) > 1: shared_ffn = self.encoder[0].ffn for blk in self.encoder[1:]: blk.ffn = shared_ffn # Pooled context summary dim depends on pool_kind. pool_dim = {"mean_last": 2 * self.D, "mean": self.D, "last": self.D}[ self.pool_kind ] # Phase-binned seasonal profile: K period-folded profiles, each # gathered by the decoder query's own phase, then mixed to D. A global # mean pool averages every phase of a cycle into one vector, so nothing # that varies with phase survives it; folding by phase keeps the # per-cycle waveform and hands each query the part of the cycle it is # forecasting. # Phase profile mixer: K periods × n_bins × (mean[,var]) → D. if self.phase_bins > 0: self.phase_mix = nn.Linear(self.K * self.stat_mult * self.D, self.D) else: self.phase_mix = None # Recency profile mixer: rb log-distance bins × (mean[,var]) # → D. Always-valid aperiodic content path. Flattened (not gathered). if self.recency_bins > 0: self.recency_mix = nn.Linear( self.recency_bins * self.stat_mult * self.D, self.D, ) else: self.recency_mix = None # Cross-cycle conv branch: a depthwise conv across cycles # at fixed phase, applied to the dominant period's [n_cycles × n_bins] # fold. Adds one D-dim feature to the query. See _cross_cycle_profile. if self.cross_cycle: self.cc_bins = self.phase_bins if self.phase_bins > 0 else 16 self.cc_cycles = 8 # most-recent N cycles folded; older clamped # Depthwise conv ACROSS the cycle axis (length cc_cycles) at fixed # phase: models how each phase evolves cycle-to-cycle. self.cc_conv = nn.Conv1d( self.D, self.D, kernel_size=3, padding=1, groups=self.D, ) self.cc_mix = nn.Linear(self.D, self.D) else: self.cc_conv = None # Decoder query input: PE + pool [+ phase D] [+ recency D] [+ cc D]. # With sig_gate, phase & recency are blended into a single D (not # concatenated), so they contribute D once, not 2·D. query_in = n_pe + pool_dim if self.sig_gate and self.phase_mix is not None and self.recency_mix is not None: query_in += self.D else: query_in += self.D if self.phase_mix is not None else 0 query_in += self.D if self.recency_mix is not None else 0 query_in += self.D if self.cross_cycle else 0 self.query_proj = nn.Linear(query_in, self.D) d_hidden = int(self.D * float(ffn_mult)) # Decoder: `decoder_depth` residual SwiGLU blocks (depth 1 is a single # block). The decoder's inputs are rich (phase/recency profiles), and # depth lets it process them. self.decoder_ffns = nn.ModuleList( [_SwiGLU(self.D, d_hidden) for _ in range(self.decoder_depth)] ) self.decoder_norms = nn.ModuleList( [nn.RMSNorm(self.D) for _ in range(self.decoder_depth)] ) # Cross-horizon coherence: a causal depthwise conv across the # horizon axis couples adjacent forecast steps (the cross-step mixing # lost when attention was dropped). Causal + fixed kernel preserves the # single-shot arbitrary-horizon property. horizon_kernel=0 disables. if self.horizon_kernel > 0: self.horizon_conv = nn.Conv1d( self.D, self.D, kernel_size=self.horizon_kernel, groups=self.D, ) self.horizon_norm = nn.RMSNorm(self.D) else: self.horizon_conv = None # Horizon-recurrent decode-state: a gated diagonal recurrence over the # SHORT horizon axis, scanning the precomputed query features. It never # re-feeds predicted values, so the decoder stays single-shot rather # than autoregressive. An unbounded carried state couples step h to ALL # earlier steps (vs the fixed-span horizon_conv), the property that # makes a recurrent decoder horizon-invariant. hr_o is zero-init so the # block is identity at start and cannot regress the baseline. if self.horizon_recurrence: self.hr_z = nn.Linear(self.D, self.D) # update gate self.hr_c = nn.Linear(self.D, self.D) # candidate self.hr_o = nn.Linear(self.D, self.D) # output proj (zero-init) nn.init.zeros_(self.hr_o.weight); nn.init.zeros_(self.hr_o.bias) self.hr_norm = nn.RMSNorm(self.D) else: self.hr_z = None self.out_proj = nn.Linear(self.D, self.n_quantiles) # Future-conv decoder (horizon-axis state evolution; the conv-native # analog of a missing-token decoder). The rest of the decoder queries a # STATIC pooled summary at every horizon position, which is why error # grows with horizon. future_conv runs a CAUSAL dilated conv over # [context-tail seed ++ seasonal-naive future fill], producing # per-future-position hidden states that EVOLVE along the horizon (each # future position is a causal-conv function of recent context + earlier # future), and injects them additively into the decoder query. The fill # is the dominant-period seasonal-naive continuation (it carries # periodic structure, so the conv evolves a real waveform forward # rather than zeros), which leaves the decoder predicting the residual # over a copy. ``fc_out`` is zero-init => EXACT baseline at start # (zero-init additive idiom), and disabling it restores that baseline. # This differs from the two cheaper readouts in the same position: a # recurrence over queries derived from the static summary adds no new # dynamics, and a phase gather only COPIES context profiles, whereas # this path RUNS the conv forward. self.future_conv = bool(future_conv) if self.future_conv: if res_adaptive: raise ValueError("future_conv is incompatible with res_adaptive") self.fc_seed = int(future_conv_seed) fc_in = 1 + n_pe # fill value + the same PE layout self.fc_in_proj = nn.Linear(fc_in, self.D) fc_dils = [2 ** i for i in range(int(future_conv_layers))] self.fc_blocks = nn.ModuleList([ _DilatedConvBlock( self.D, kernel=int(kernel), dilation=int(d_i), ffn_mult=float(ffn_mult), causal=True, separable=True, # auxiliary module: keep it light (~51K add) ) for d_i in fc_dils ]) # Weight-tie the FFN across fc blocks: the convs carry the horizon # dynamics, and one shared FFN keeps the param add modest. shared = self.fc_blocks[0].ffn for blk in self.fc_blocks[1:]: blk.ffn = shared self.fc_out = nn.Linear(self.D, self.D) # zero-init => exact baseline nn.init.zeros_(self.fc_out.weight) nn.init.zeros_(self.fc_out.bias) self.base_seasonality = float(base_seasonality) # ---- helpers ---------------------------------------------------------- def _future_conv_states( self, h: torch.Tensor, fut_pe: torch.Tensor, fill: torch.Tensor, ) -> torch.Tensor: """Causal-conv continuation states at the H future positions. h: (B, L, D) encoder output; fut_pe: (B, H, n_pe); fill: (B, H) the seasonal-naive future continuation. Returns (B, H, D). """ # self.L is static (asserted == L in forward), so using it instead of # the traced h.shape[1] keeps dynamo from graph-splitting on a symint # bound. ft = self.fc_in_proj(torch.cat([fill.unsqueeze(-1), fut_pe], dim=-1)) # (B,H,D) seed = h[:, -min(self.fc_seed, self.L):, :] # (B, seed, D) z = torch.cat([seed, ft], dim=1) # (B, seed+H, D) for blk in self.fc_blocks: # causal: no future leak z = blk(z) return z[:, -ft.shape[1]:, :] # (B, H, D) @torch.compiler.disable() def _detect_periods( self, x_fp32: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Run the significance-filtered periodogram in fp32 (no grad). @torch.compiler.disable: the periodogram uses a complex rfft that Inductor cannot codegen: left inside the compiled graph it forces an eager fallback + graph break every forward, early in `forward`, blocking fusion of the whole conv encoder/decoder downstream. Disabling compile on this (no_grad, fp32, produces integer periods the rest only reads) makes a clean eager boundary: the FFT runs eager, everything after fuses. Output bit-identical (only where it compiles changes). Returns ``(periods, n_valid, scores)``: the integer periods (0 = rejected), the count of significant periods per sample, and the per-period periodogram scores. ``n_valid`` is the periodicity-strength signal the significance gate reads. """ B, L = x_fp32.shape if self.periodogram_off: # Control: no period detection → phase encoding zeros out, phase # machinery is inert. Tests whether the conv backbone matches with # a phase-free (recency-only) decoder. z = torch.zeros(B, self.K, dtype=torch.long, device=x_fp32.device) return (z, torch.zeros(B, dtype=torch.long, device=x_fp32.device), torch.zeros(B, self.K, device=x_fp32.device)) with torch.no_grad(): periods, scores, n_valid = significant_periods( x_fp32, min_period=2, max_period=L // 2, top_k=self.K, significance_alpha=self.significance_alpha, ) periods = periods.long() if self.min_cycles > 0: # "Do no harm": only TRUST a period with >= min_cycles full # cycles in the window (period <= L/min_cycles). Periods too long # to be reliably estimated (e.g. an 8640-sample daily cycle in a # 2048 window) are zeroed → the significance gate routes those # series to the recency/local path instead of mis-locking phase. max_p = L // self.min_cycles keep = (periods > 0) & (periods <= max_p) periods = torch.where(keep, periods, torch.zeros_like(periods)) n_valid = keep.sum(dim=1).long() return periods, n_valid.long(), scores.float() def _pool(self, h: torch.Tensor) -> torch.Tensor: """Pool encoder output to a fixed-size summary. h: (B, L, D) Returns: (B, pool_dim) """ if self.pool_kind == "mean": return h.mean(dim=1) if self.pool_kind == "last": return h[:, -1, :] # mean_last: return torch.cat([h.mean(dim=1), h[:, -1, :]], dim=-1) def _scatter_profile( self, h: torch.Tensor, bins: torch.Tensor, nb: int, weight: torch.Tensor | None = None, ) -> torch.Tensor: """Weighted scatter-mean (+ optional per-bin variance) of h into nb bins. h: (B, L, D) bins: (B, L) int in [0, nb) weight: (B, L) non-negative, or None for uniform. Returns (B, nb, stat_mult·D): per-bin mean, then per-bin variance if ``phase_stats == 'mean_var'``. Empty bins → global mean (and zero variance). """ B, L, D = h.shape oh = F.one_hot(bins, nb).to(h.dtype) # (B,L,nb) if weight is not None: oh = oh * weight.unsqueeze(-1) cnt = oh.sum(dim=1).unsqueeze(-1) # (B,nb,1) denom = cnt.clamp(min=1e-6) mean = torch.bmm(oh.transpose(1, 2), h) / denom # (B,nb,D) gmean = h.mean(dim=1, keepdim=True) # (B,1,D) empty = cnt <= 0 mean = torch.where(empty, gmean.expand(B, nb, D), mean) if self.phase_stats == "mean_var": sq = torch.bmm(oh.transpose(1, 2), h * h) / denom # E[h²] var = (sq - mean * mean).clamp(min=0.0) var = torch.where(empty, torch.zeros_like(var), var) return torch.cat([mean, var], dim=-1) # (B,nb,2D) return mean # (B,nb,D) def _recency_weight(self, L: int, device) -> torch.Tensor | None: """exp recency weight over context positions: recent positions weigh more, so the folded profile tracks the CURRENT regime's waveform rather than the window average. None if tau<=0 (uniform).""" if self.phase_recency_tau <= 0.0: return None t = torch.arange(L, device=device).float() dist = (L - 1 - t) / L # 0 at now return torch.exp(-dist / self.phase_recency_tau).view(1, L) def _phase_profile( self, h: torch.Tensor, periods: torch.Tensor, ) -> torch.Tensor: """Period-fold the encoder output into per-phase profiles. Returns (B, K, n_bins, stat_mult·D). Each detected period p_k folds the sequence into n_bins phase bins; per bin we keep the (recency- weighted) mean [and variance]. Empty bins → global mean. """ B, L, D = h.shape K, nb = self.K, self.phase_bins t = torch.arange(L, device=h.device).view(1, L).float() # (1,L) p_safe = periods.clamp(min=1).float() # (B,K) weight = self._recency_weight(L, h.device) if weight is not None: weight = weight.expand(B, L) profs = [] for k in range(K): frac = (t % p_safe[:, k:k + 1]) / p_safe[:, k:k + 1] # (B,L) bins = torch.clamp((frac * nb).long(), max=nb - 1) # (B,L) profs.append(self._scatter_profile(h, bins, nb, weight)) return torch.stack(profs, dim=1) # (B,K,nb,S·D) def _gather_phase( self, prof: torch.Tensor, fut_pos: torch.Tensor, periods: torch.Tensor, weight: torch.Tensor | None = None, ) -> torch.Tensor: """Gather each future query's matching phase bin per period, mix → D. prof: (B, K, n_bins, S·D); fut_pos: (B, H); periods: (B, K). weight: optional (B,K) per-period reliability weight (period_trust). Returns (B, H, D). """ B, K, nb, SD = prof.shape H = fut_pos.shape[1] p_safe = periods.clamp(min=1).float().view(B, 1, K) # (B,1,K) frac = (fut_pos.unsqueeze(-1).float() % p_safe) / p_safe # (B,H,K) fbins = torch.clamp((frac * nb).long(), max=nb - 1) # (B,H,K) idx = fbins.permute(0, 2, 1).unsqueeze(-1).expand(B, K, H, SD) gathered = torch.gather(prof, 2, idx) # (B,K,H,S·D) if weight is not None: gathered = gathered * weight.view(B, K, 1, 1).to(gathered.dtype) gathered = gathered.permute(0, 2, 1, 3).reshape(B, H, K * SD) return self.phase_mix(gathered) # (B,H,D) def _period_trust_weights( self, periods: torch.Tensor, scores: torch.Tensor, ) -> torch.Tensor: """Hyperparameter-free per-period reliability weight w_k∈[0,1] (B,K). ln-coverage = ln(L/period) (data/structure-determined); for 'full' also the significance margin ln(s_k/t_alpha) (s_k = the periodogram score, t_alpha = data-determined Bonferroni threshold). The sigmoid crossover is LEARNED (the linear layer's weights and bias), not a hand-set threshold. 0 on rejected slots. """ valid = periods > 0 logcov = torch.log(float(self.L) / periods.clamp(min=1).float()) # (B,K) if self.period_trust == "coverage": feat = logcov.unsqueeze(-1) # (B,K,1) else: margin = torch.log(scores.clamp(min=1e-12) / self._t_alpha) # >=0 for survivors feat = torch.stack([margin, logcov], dim=-1) # (B,K,2) w = torch.sigmoid(self.pt(feat.to(self.pt.weight.dtype))).squeeze(-1) return torch.where(valid, w, torch.zeros_like(w)) def _recency_feat(self, h: torch.Tensor) -> torch.Tensor: """Recency-binned profile: bin context positions by log-distance-from-now and pool. Always valid (no period needed); the aperiodic content path. Flattened to a single (B, D) descriptor (broadcast to all horizons; the query's own PE carries how-far-ahead). """ B, L, D = h.shape rb = self.recency_bins t = torch.arange(L, device=h.device).view(1, L).float().expand(B, L) dist = (L - 1 - t).clamp(min=0.0) # 0=now frac = torch.log1p(dist) / math.log1p(float(L - 1) + 1e-9) bins = torch.clamp((frac * rb).long(), max=rb - 1) # (B,L) prof = self._scatter_profile(h, bins, rb, None) # (B,rb,S·D) return self.recency_mix(prof.reshape(B, rb * self.stat_mult * D)) def _cross_cycle_profile( self, h: torch.Tensor, periods: torch.Tensor, ) -> torch.Tensor: """Cross-cycle conv, true ragged form. For the dominant period p0, fold the sequence into a [cycles-back-from-now × phase] grid (B, nc, nb, D) by scatter-mean, then convolve ACROSS the cycle axis at fixed phase (depthwise conv1d over nc), modelling how each phase evolves cycle-to-cycle ("every Monday 9am, trending up"). Read out the most-recent cycle (post-conv, so it has seen the trend). Returns a (B, nb, D) phase profile the decoder gathers by its own phase. Per-sample period handled like phase-binning: phase resampled to nb fixed bins; cycles-back clamped to nc (older cycles fold into the oldest slot). Attention-free, fixed-shape, batchable. """ B, L, D = h.shape nb, nc = self.cc_bins, self.cc_cycles t = torch.arange(L, device=h.device).view(1, L).float() # (1,L) p0 = periods[:, :1].clamp(min=1).float() # (B,1) pbin = torch.clamp(((t % p0) / p0 * nb).long(), max=nb - 1) # (B,L) cyc = torch.clamp(((L - 1 - t) // p0).long(), max=nc - 1) # (B,L) 0=now comb = (cyc * nb + pbin).clamp(min=0, max=nc * nb - 1) # (B,L) oh = F.one_hot(comb, nc * nb).to(h.dtype) # (B,L,nc·nb) cnt = oh.sum(dim=1).unsqueeze(-1).clamp(min=1e-6) grid = (torch.bmm(oh.transpose(1, 2), h) / cnt).view(B, nc, nb, D) # conv across cycles (nc) at fixed phase, per channel. x = grid.permute(0, 2, 3, 1).reshape(B * nb, D, nc) # (B·nb, D, nc) x = self.cc_conv(x).reshape(B, nb, D, nc) return x[..., 0] # most-recent cycle (B,nb,D) @staticmethod def _prefix_integral( f: torch.Tensor, Csum: torch.Tensor, xpad: torch.Tensor, L: int, ) -> torch.Tensor: """Integral of piecewise-constant x from 0 to fractional position f. f: (B,M) in native units. Csum: (B,L+1) prefix sums; xpad: (B,L+1).""" fc = f.clamp(0.0, float(L)) k = torch.floor(fc).long() rem = (fc - k.float()).to(Csum.dtype) return torch.gather(Csum, 1, k) + rem * torch.gather(xpad, 1, k) def _resolution_adapt( self, x: torch.Tensor, periods: torch.Tensor, H: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Resolution adaptation (the TinyCast premise; zero params). Resample the context onto a canonical-period grid so the fixed dilation schedule spans consistent CYCLE-fractions across all sampling rates: the dominant detected period is warped to ``res_period_target`` samples/cycle. All detected periods scale by the same ratio; the future native horizon is queried at its canonical-mapped position (decoder is position- parameterized, so outputs are native values and no output resampling is needed). Returns (x_canon (B,L), periods_canon (B,K), fut_pos_canon (B,H) float). Aperiodic series (dominant period 0) pass through unchanged (r=1). """ B, L = x.shape pt = float(self.res_period_target) p0 = periods[:, 0].float() # (B,) dominant r = torch.where(p0 > 0, pt / p0.clamp(min=1.0), torch.ones_like(p0)) # res_r_max=1.0 → downsample-only (high-freq squeezed to canonical; # low-freq left at native, no history truncation / no Δ blow-up). r = r.clamp(0.05, self.res_r_max).view(B, 1) # canonical per native j = torch.arange(L, device=x.device).view(1, L).float() # canonical idx 0..L-1 # native time (center) for canonical index j; now = most recent canonical. t = (L - 1) - ((L - 1) - j) / r # (B,L), <0 = pre-context # Upsample/identity (r>=1): linear interpolation (true pass-through at # r=1). Downsample (r<1): area-average over the native window w=1/r, # ANTI-ALIASED (averages the whole window, not 2 endpoints). Zero params. t0 = torch.floor(t) frac = (t - t0).to(x.dtype) x_lin = (torch.gather(x, 1, t0.clamp(0, L - 1).long()) * (1.0 - frac) + torch.gather(x, 1, (t0 + 1).clamp(0, L - 1).long()) * frac) w = 1.0 / r # (B,1) native window width lo, hi = t - w / 2.0, t + w / 2.0 Csum = F.pad(x.cumsum(dim=1), (1, 0)) # (B,L+1): Csum[k]=Σ x[:k] xpad = F.pad(x, (0, 1)) # (B,L+1): x[L]=0 denom = (hi.clamp(0.0, L) - lo.clamp(0.0, L)).clamp(min=1e-6) x_area = (self._prefix_integral(hi, Csum, xpad, L) - self._prefix_integral(lo, Csum, xpad, L)) / denom x_canon = torch.where(r < 1.0, x_area, x_lin) # anti-alias only on downsample x_canon = x_canon * (t >= 0).to(x.dtype) # mask pre-context → 0 periods_canon = torch.round(periods.float() * r).long() periods_canon = torch.where( periods > 0, periods_canon.clamp(min=2), torch.zeros_like(periods), ) h_steps = torch.arange(1, H + 1, device=x.device).view(1, H).float() fut_pos_canon = (L - 1) + h_steps * r # (B,H) canonical return x_canon, periods_canon, fut_pos_canon def _seasonal_naive( self, x: torch.Tensor, fut_pos: torch.Tensor, periods: torch.Tensor, ) -> torch.Tensor: """Value-space seasonal-naive baseline (zero params). Fold the (normalized) input x by the dominant period into phase bins, take the per-phase mean VALUE, and gather the bin matching each future query's phase. The network then learns only the residual on top of this baseline: a target reframe, not added model complexity. Aperiodic / empty-bin → fall back to the context mean (persistence-of-level). x: (B, L) fut_pos: (B, H) periods: (B, K) → (B, H) """ B, L = x.shape nb = self.phase_bins if self.phase_bins > 0 else 16 t = torch.arange(L, device=x.device).view(1, L).float() p0 = periods[:, :1].clamp(min=1).float() # (B,1) dominant pbin = torch.clamp(((t % p0) / p0 * nb).long(), max=nb - 1) # (B,L) oh = F.one_hot(pbin, nb).to(x.dtype) # (B,L,nb) cnt = oh.sum(dim=1) # (B,nb) base = torch.bmm(oh.transpose(1, 2), x.unsqueeze(-1)).squeeze(-1) # (B,nb) gmean = x.mean(dim=1, keepdim=True) # (B,1) base = torch.where(cnt > 0, base / cnt.clamp(min=1.0), gmean.expand(B, nb)) fb = torch.clamp((fut_pos.float() % p0) / p0 * nb, max=nb - 1).long() # (B,H) return torch.gather(base, 1, fb) # (B,H) def _super_naive( self, x: torch.Tensor, fut_pos: torch.Tensor, periods: torch.Tensor, ) -> torch.Tensor: """Multi-seasonal "super-naive" baseline (zero params). Greedy additive decomposition over ALL significant periods: start from the context mean, then for each significant period (strongest first) fold the running residual into per-phase means, subtract it (deflate), and accumulate that component's value at the future phase. Result: baseline(L+h) = mean + Σ_k s_k[phase_k(L+h)], the genuine multi-period seasonal-naive forecast. Non-significant periods (p=0) contribute zero. x: (B, L) fut_pos: (B, H) periods: (B, K) → (B, H) """ B, L = x.shape H = fut_pos.shape[1] nb = self.phase_bins if self.phase_bins > 0 else 16 t = torch.arange(L, device=x.device).view(1, L).float() if self.residual_trend: # Level term = linear trend (closed-form LS), extrapolated forward. # baseline = trend + seasonal, the classical decomposition. tc = t - t.mean() # centered (1,L) xc = x - x.mean(dim=1, keepdim=True) # (B,L) slope = (tc * xc).sum(1, keepdim=True) / (tc * tc).sum().clamp(min=1.0) intercept = x.mean(dim=1, keepdim=True) # value at centered t=0 tmean = t.mean() trend_ctx = intercept + slope * (t - tmean) # (B,L) r = x - trend_ctx # de-trended residual baseline = intercept + slope * (fut_pos.float() - tmean) # (B,H) trend extrap else: mean = x.mean(dim=1, keepdim=True) # (B,1) r = x - mean # residual baseline = mean.expand(B, H).clone() # (B,H) for k in range(self.K): pk = periods[:, k:k + 1].float() # (B,1), 0 if not sig sig = (pk > 0).to(x.dtype) # (B,1) pks = pk.clamp(min=1.0) pbin = torch.clamp((t % pks) / pks * nb, max=nb - 1).long() # (B,L) oh = F.one_hot(pbin, nb).to(x.dtype) # (B,L,nb) cnt = oh.sum(dim=1).clamp(min=1.0) # (B,nb) s_k = torch.bmm(oh.transpose(1, 2), r.unsqueeze(-1)).squeeze(-1) / cnt s_k = s_k * sig # (B,nb), zero if not sig r = r - torch.gather(s_k, 1, pbin) # deflate fb = torch.clamp((fut_pos.float() % pks) / pks * nb, max=nb - 1).long() baseline = baseline + torch.gather(s_k, 1, fb) # (B,H) return baseline def _gather_cc( self, cc_prof: torch.Tensor, fut_pos: torch.Tensor, periods: torch.Tensor, ) -> torch.Tensor: """Gather each future query's matching phase bin from the cross-cycle profile (dominant period), mix → D. cc_prof: (B,nb,D).""" B, nb, D = cc_prof.shape H = fut_pos.shape[1] p0 = periods[:, :1].clamp(min=1).float() # (B,1) fb = torch.clamp((fut_pos.float() % p0) / p0 * nb, max=nb - 1).long() gathered = torch.gather(cc_prof, 1, fb.unsqueeze(-1).expand(B, H, D)) return self.cc_mix(gathered) # (B,H,D) # ---- forward ---------------------------------------------------------- def _local_anchor_channels( self, x: torch.Tensor, scale_factor: torch.Tensor | float | None, ) -> torch.Tensor: """Two causal local-statistics channels exposing the recent level/scale to the encoder (the gradient-connected re-anchoring signal WindowMinMax lacks): ch1 = (x_t - m_t) / (s_t + eps) local-scale residual (a causal z-score) ch2 = log(s_t + eps) log local scale (global normed range ~= 1) m_t, s_t = causal boxcar mean / std over a trailing window w ~ one canonical period round(base_seasonality / scale_factor), clamped [8, L//4], fallback 64. Vectorized via cumsum + per-sample-window gather (O(L), no python loop). """ B, L = x.shape device = x.device if scale_factor is not None: sf = (scale_factor if torch.is_tensor(scale_factor) else x.new_tensor(scale_factor)).reshape(-1).float() if sf.numel() == 1: sf = sf.expand(B) w = (self.base_seasonality / sf.clamp(min=1e-3)).round().long() w = w.clamp(min=8, max=max(8, L // 4)) else: w = torch.full((B,), 64, device=device, dtype=torch.long) # Center by the per-series mean before a FP32 cumsum. The two-pass variance # (E[x^2]-E[x]^2) over a length-L cumsum otherwise suffers catastrophic # cancellation on long flat/sparse regions; variance is shift-invariant, so # centering changes nothing but keeps the cumsum magnitudes small enough that # fp32 stays accurate there, at no extra memory. xf = x.float() xc = xf - xf.mean(dim=1, keepdim=True) cs = F.pad(torch.cumsum(xc, dim=1), (1, 0)) # (B, L+1), cs[:,0]=0 cs2 = F.pad(torch.cumsum(xc * xc, dim=1), (1, 0)) t = torch.arange(L, device=device).view(1, L).expand(B, L) lo = (t - w.view(B, 1) + 1).clamp(min=0) # trailing-window start cnt = (t - lo + 1).float() # window length (>= 1) sum_x = cs.gather(1, t + 1) - cs.gather(1, lo) sum_x2 = cs2.gather(1, t + 1) - cs2.gather(1, lo) m = sum_x / cnt # centered local mean s = (sum_x2 / cnt - m * m).clamp(min=0.0).sqrt() # local std (shift-invariant) eps = 1e-4 # floor vs the unit normed range -> flat regions give ch1 ~ 0, no blowup ch1 = (xc - m) / (s + eps) # = (x - local mean)/(s+eps) ch2 = torch.log(s + eps) return torch.stack([ch1, ch2], dim=-1).to(x.dtype) # (B, L, 2) def forward( self, x_normed: torch.Tensor, nan_mask: torch.Tensor | None = None, scale_factor: torch.Tensor | float | None = None, horizon: int | None = None, ) -> torch.Tensor: # observed-mask (1=observed, 0=missing) for the missing-value channel. # res_adaptive is rejected with missing_channel (see __init__), so this # mask stays aligned with x throughout. obs_mask = None if self.missing_channel and nan_mask is not None: obs_mask = nan_mask[..., 0] if nan_mask.dim() == 3 else nan_mask # (B,L) if x_normed.dim() == 3 and x_normed.shape[-1] > 1: x = x_normed[..., 0] elif x_normed.dim() == 3: x = x_normed.squeeze(-1) else: x = x_normed # (B, L) x = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0) B, L = x.shape assert L == self.L, f"context length mismatch: got {L}, expected {self.L}" H = self.p_out if horizon is None else int(horizon) device = x.device # Period detection in fp32 (no grad). periods, n_valid, scores = self._detect_periods(x.float()) # (B,K),(B,),(B,K) # Resolution adaptation: warp context to a canonical cycle-resolution so # the fixed dilations span consistent cycle-fractions across rates. if self.res_adaptive: x, periods, fut_pos = self._resolution_adapt(x, periods, H) ctx_pos = torch.arange(L, device=device).view(1, L).expand(B, L) else: ctx_pos = torch.arange(L, device=device).view(1, L).expand(B, L) fut_pos = torch.arange(L, L + H, device=device).view(1, H).expand(B, H) with torch.amp.autocast( device_type=device.type if x.is_cuda else "cpu", enabled=False, ): ctx_pe = _positional_encoding( ctx_pos, periods, L, n_harmonics=self.n_harmonics, ) fut_pe = _positional_encoding( fut_pos, periods, L, n_harmonics=self.n_harmonics, ) ctx_pe = ctx_pe.to(x.dtype) fut_pe = fut_pe.to(x.dtype) # Per-period reliability weight (period_trust; off in the released config). # Down-weight unreliable/spurious periods continuously. Applied to the # phase-encoding channels here, and to the phase-binning gather + gate below. ptw = None if self.pt is not None: ptw = self._period_trust_weights(periods, scores).to(x.dtype) # (B,K) # phase channels are the first n_phase cols, laid out per period as # n_harmonics*2 consecutive channels → repeat each w_k that many times. rep = self.n_harmonics * 2 chan_w = ptw.repeat_interleave(rep, dim=1).view(B, 1, -1) # (B,1,n_phase) np_ = chan_w.shape[-1] ctx_pe = torch.cat([ctx_pe[..., :np_] * chan_w, ctx_pe[..., np_:]], dim=-1) fut_pe = torch.cat([fut_pe[..., :np_] * chan_w, fut_pe[..., np_:]], dim=-1) # Embed context. if self.decompose_kernel > 0: # moving-average series decomposition: moving-average trend + # seasonal residual, fed as two channels. k = self.decompose_kernel pad = k // 2 xp = F.pad(x.unsqueeze(1), (pad, pad), mode="replicate") # (B,1,L+2pad) trend = F.avg_pool1d(xp, kernel_size=k, stride=1)[..., :L].squeeze(1) seasonal = x - trend value_ch = torch.stack([trend, seasonal], dim=-1) # (B,L,2) else: value_ch = x.unsqueeze(-1) # (B,L,1) if self.missing_channel and obs_mask is not None: value_ch = torch.cat( [value_ch, obs_mask.unsqueeze(-1).to(value_ch.dtype)], dim=-1 ) # (B,L,nv+1) if self.local_anchor: value_ch = torch.cat( [value_ch, self._local_anchor_channels(x, scale_factor)], dim=-1 ) # (B,L,nv+2) ctx_in = torch.cat([value_ch, ctx_pe], dim=-1) # (B,L,nv+n_pe) h = self.in_proj(ctx_in) # (B, L, D) # Dilated-conv encoder. for block in self.encoder: h = block(h) # Per-horizon context feature: static pooled summary, broadcast to all H. summary = self._pool(h) # (B, pool_dim) ctx_feat = summary.unsqueeze(1).expand(B, H, -1) # (B, H, pool_dim) # Decoder: per-horizon query from PE(L+h) + context feature # [+ phase profile] [+ recency profile] [+ cross-cycle feat]. q_parts = [fut_pe, ctx_feat] phase_feat = None if self.phase_mix is not None: prof = self._phase_profile(h, periods) # (B,K,nb,S·D) phase_feat = self._gather_phase(prof, fut_pos, periods, weight=ptw) # (B,H,D) rec_feat = None if self.recency_mix is not None: rec_feat = self._recency_feat(h).unsqueeze(1).expand(B, H, self.D) if self.sig_gate and phase_feat is not None and rec_feat is not None: # Blend by periodicity strength: many significant periods → trust # the phase profile; few/none → lean on the recency profile. With # period_trust, use the soft Σw instead of the integer n_valid. strength = ptw.sum(dim=1) if ptw is not None else n_valid.float() g = (strength / float(self.K)).clamp(0.0, 1.0).view(B, 1, 1) q_parts.append(g * phase_feat + (1.0 - g) * rec_feat) else: if phase_feat is not None: q_parts.append(phase_feat) if rec_feat is not None: q_parts.append(rec_feat) if self.cross_cycle: cc_prof = self._cross_cycle_profile(h, periods) # (B,nb,D) q_parts.append(self._gather_cc(cc_prof, fut_pos, periods)) # (B,H,D) q_in = torch.cat(q_parts, dim=-1) q = self.query_proj(q_in) # (B, H, D) if self.future_conv: # Horizon-axis evolving states from a causal conv over context-tail # + seasonal-naive fill, injected additively (fc_out zero-init => # exact baseline at init). Gives the static-summary decoder the # per-horizon dynamics it lacks. fill = self._seasonal_naive(x, fut_pos, periods) # (B, H) fut_states = self._future_conv_states(h, fut_pe, fill) # (B, H, D) q = q + self.fc_out(fut_states) for ffn, norm in zip(self.decoder_ffns, self.decoder_norms): q = _norm_fp32(norm, q + ffn(q)) if self.horizon_conv is not None: # Causal conv over the horizon axis: pad (k-1) on the left so step # h sees only h, h-1, …, h-(k-1): no future leakage, any H. qt = F.pad(q.transpose(1, 2), (self.horizon_kernel - 1, 0)) hc = self.horizon_conv(qt).transpose(1, 2) # (B,H,D) q = _norm_fp32(self.horizon_norm, q + hc) if self.hr_z is not None: # gated-recurrence decode-state: gated diagonal recurrence over the H axis. # Scans the query FEATURES only (no value feedback). Sequential # scan: the horizon is short, so this is cheap and stable. z = torch.sigmoid(self.hr_z(q)) # (B,H,D) update gate c = torch.tanh(self.hr_c(q)) # (B,H,D) candidate s = torch.zeros(B, self.D, dtype=q.dtype, device=q.device) states = [] for t in range(q.shape[1]): s = (1.0 - z[:, t]) * s + z[:, t] * c[:, t] states.append(s) hstate = torch.stack(states, dim=1) # (B,H,D) q = _norm_fp32(self.hr_norm, q + self.hr_o(hstate)) # hr_o zero-init y = self.out_proj(q) # (B, H, Q) if self.residual_naive: # Learn the residual over a (multi-)seasonal-naive baseline. if self.residual_multi: baseline = self._super_naive(x, fut_pos, periods) # (B, H) else: baseline = self._seasonal_naive(x, fut_pos, periods) # (B, H) y = y + baseline.unsqueeze(-1) return y