"""Continuous-horizon sequence models for forecasting-v4.""" from __future__ import annotations import math from collections.abc import Sequence import torch from torch import nn class PooledLinear(nn.Module): def __init__( self, *, context_length: int, output_horizons: Sequence[float], classes: int, bins: int = 8, ) -> None: super().__init__() if context_length % bins: raise ValueError("context length must be divisible by pool bins") self.bins = bins self.horizon_count = len(output_horizons) self.classes = classes outputs = self.horizon_count * classes self.output = nn.Linear(bins * 5, outputs) self.estimated_madds = ( context_length * 4 + bins * 5 * outputs ) def forward( self, tokens: torch.Tensor, padding_mask: torch.Tensor, horizons: torch.Tensor | None = None, ) -> torch.Tensor: del padding_mask if horizons is not None: raise ValueError("fixed pooled control cannot query horizons") if tokens.shape[2] != 4: raise ValueError("pooled control requires baseline plus padding") batch, length, _channels = tokens.shape segment = tokens.reshape( batch, self.bins, length // self.bins, tokens.shape[2], ) observed = segment[..., 1] valid = segment[..., 3] observed_count = observed.sum(dim=2).clamp_min(1.0) valid_count = valid.sum(dim=2).clamp_min(1.0) return_mean = (segment[..., 0] * observed).sum(dim=2) / observed_count return_rms = torch.sqrt( (segment[..., 0].square() * observed).sum(dim=2) / observed_count + 1e-8 ) state = torch.stack( ( return_mean, return_rms, observed.mean(dim=2), (segment[..., 2] * valid).sum(dim=2) / valid_count, valid.mean(dim=2), ), dim=2, ).reshape(batch, -1) return self.output(state).reshape( batch, self.horizon_count, self.classes, ) class CausalResidualBlock(nn.Module): def __init__(self, width: int, dilation: int) -> None: super().__init__() self.padding = 2 * dilation self.convolution = nn.Conv1d( width, width, kernel_size=3, dilation=dilation, padding=self.padding, ) self.normalization = nn.GroupNorm(1, width) self.activation = nn.GELU() def forward(self, values: torch.Tensor) -> torch.Tensor: residual = values values = self.convolution(values) values = values[:, :, : -self.padding] return self.activation(self.normalization(values)) + residual class CausalTCN(nn.Module): def __init__( self, *, context_length: int, channels: int, output_horizons: Sequence[float], classes: int, target_parameters: int, ) -> None: super().__init__() blocks = max(1, math.ceil(math.log2(context_length))) outputs = len(output_horizons) * classes width = 8 for candidate in range(8, 513): estimate = ( channels * candidate + candidate + blocks * (3 * candidate * candidate + 3 * candidate) + candidate * outputs + outputs ) if estimate > target_parameters: break width = candidate self.width = width self.block_count = blocks self.input_projection = nn.Conv1d(channels, width, kernel_size=1) self.blocks = nn.Sequential( *[ CausalResidualBlock(width, 2**index) for index in range(blocks) ] ) self.output = nn.Linear(width, outputs) self.horizon_count = len(output_horizons) self.classes = classes self.estimated_madds = ( context_length * channels * width + blocks * context_length * 3 * width * width + width * outputs ) def forward( self, tokens: torch.Tensor, padding_mask: torch.Tensor, horizons: torch.Tensor | None = None, ) -> torch.Tensor: del padding_mask if horizons is not None: raise ValueError("fixed-head TCN cannot query custom horizons") values = self.input_projection(tokens.transpose(1, 2)) values = self.blocks(values) result = self.output(values[:, :, -1]) return result.reshape(-1, self.horizon_count, self.classes) class ContinuousHorizonTransformer(nn.Module): def __init__( self, *, context_length: int, channels: int, output_horizons: Sequence[float], classes: int, target_parameters: int, query_mode: str, rotary_base: float, ) -> None: super().__init__() if query_mode not in ("learned", "scalar_mlp", "rotary"): raise ValueError("invalid horizon query mode") if rotary_base <= 1.0: raise ValueError("rotary base must exceed one") output_horizons = tuple(float(value) for value in output_horizons) if any(value <= 0.0 for value in output_horizons): raise ValueError("output horizons must be positive") width = 16 closest = float("inf") for candidate in range(16, 257, 4): estimate = ( 20 * candidate * candidate + candidate * ( context_length + channels + len(output_horizons) + classes + 27 ) + classes ) distance = abs(estimate - target_parameters) if distance < closest: width = candidate closest = distance if width % 2: raise ValueError("continuous query width must be even") heads = 4 if width % 4 == 0 else 2 self.width = width self.query_mode = query_mode self.rotary_base = float(rotary_base) self.classes = classes self.input_projection = nn.Linear(channels, width) self.position = nn.Parameter(torch.empty(context_length, width)) layer = nn.TransformerEncoderLayer( d_model=width, nhead=heads, dim_feedforward=width * 2, dropout=0.0, activation="gelu", batch_first=True, norm_first=True, ) self.encoder = nn.TransformerEncoder(layer, num_layers=2) self.query_attention = nn.MultiheadAttention( width, heads, dropout=0.0, batch_first=True, ) self.output = nn.Linear(width, classes) self.register_buffer( "output_horizons", torch.tensor(output_horizons, dtype=torch.float32), persistent=True, ) if query_mode == "learned": self.learned_queries = nn.Parameter( torch.empty(len(output_horizons), width) ) elif query_mode == "scalar_mlp": scalar_width = max(4, width // 4) self.scalar_query = nn.Sequential( nn.Linear(1, scalar_width), nn.GELU(), nn.Linear(scalar_width, width), ) else: self.rotary_query = nn.Parameter(torch.empty(width)) self.linear_query = nn.Parameter(torch.empty(width)) pair = torch.arange(width // 2, dtype=torch.float32) self.register_buffer( "rotary_frequency", self.rotary_base ** (-2.0 * pair / width), persistent=True, ) nn.init.normal_(self.position, std=0.02) if query_mode == "learned": nn.init.normal_(self.learned_queries, std=0.02) elif query_mode == "rotary": nn.init.normal_(self.rotary_query, std=0.02) nn.init.normal_(self.linear_query, std=0.02) horizon_count = len(output_horizons) encoder_madds = 2 * ( 8 * context_length * width * width + 2 * context_length * context_length * width ) query_madds = ( 2 * (horizon_count + context_length) * width * width + 2 * horizon_count * context_length * width ) self.estimated_madds = ( context_length * channels * width + encoder_madds + query_madds + horizon_count * width * classes ) def _queries(self, horizons: torch.Tensor) -> torch.Tensor: values = horizons.to( device=self.position.device, dtype=self.position.dtype, ) coordinate = torch.log2(values / 2.0) if self.query_mode == "learned": if ( len(values) != len(self.output_horizons) or not torch.allclose( values, self.output_horizons.to(values), ) ): raise ValueError( "learned queries only support configured horizons" ) return self.learned_queries if self.query_mode == "scalar_mlp": return self.scalar_query(coordinate[:, None]) angles = coordinate[:, None] * self.rotary_frequency[None, :] base = self.rotary_query.reshape(-1, 2) even = base[:, 0][None, :] odd = base[:, 1][None, :] cosine = torch.cos(angles) sine = torch.sin(angles) rotated_even = even * cosine - odd * sine rotated_odd = even * sine + odd * cosine rotated = torch.stack( (rotated_even, rotated_odd), dim=2, ).reshape(len(values), self.width) return rotated + coordinate[:, None] * self.linear_query[None, :] def forward( self, tokens: torch.Tensor, padding_mask: torch.Tensor, horizons: torch.Tensor | None = None, ) -> torch.Tensor: values = self.input_projection(tokens) + self.position[None, :, :] encoded = self.encoder( values, src_key_padding_mask=padding_mask, ) query_horizons = ( self.output_horizons if horizons is None else horizons ) queries = self._queries(query_horizons) queries = queries[None, :, :].expand(len(tokens), -1, -1) decoded, _weights = self.query_attention( queries, encoded, encoded, key_padding_mask=padding_mask, need_weights=False, ) return self.output(decoded) def build_model( name: str, *, context_length: int, channels: int, output_horizons: Sequence[float], classes: int, target_parameters: int, rotary_base: float, ) -> nn.Module: if name == "pooled_linear": return PooledLinear( context_length=context_length, output_horizons=output_horizons, classes=classes, ) if name == "tcn": return CausalTCN( context_length=context_length, channels=channels, output_horizons=output_horizons, classes=classes, target_parameters=target_parameters, ) query_modes = { "transformer_learned": "learned", "transformer_scalar": "scalar_mlp", "transformer_rotary": "rotary", } if name in query_modes: return ContinuousHorizonTransformer( context_length=context_length, channels=channels, output_horizons=output_horizons, classes=classes, target_parameters=target_parameters, query_mode=query_modes[name], rotary_base=rotary_base, ) raise ValueError(f"unknown forecasting-v4 model: {name}")