| """ |
| A small, patch-based transformer for zero-shot time series forecasting — |
| a toy-scale version of the same general recipe used by real time-series |
| foundation models (Chronos, TimesFM, t0-alpha, etc.): |
| |
| 1. Per-series normalization: each input window is normalized by its |
| OWN mean and std before entering the model, and de-normalized |
| using those same stats afterward. This is the single most |
| important design choice for generalization — it means the model |
| learns SHAPE (trend, seasonality, noise structure), not absolute |
| numbers, so a model trained on synthetic series scaled 0.5-500 can |
| meaningfully forecast a real series scaled in the thousands. |
| |
| 2. Patching: rather than feeding every single timestep as a separate |
| token (expensive, and treats each point independently), the input |
| window is split into small chunks ("patches") and each patch is |
| projected into the model's embedding space as one token. This is |
| the same idea PatchTST and most modern time-series transformers |
| use — it's both more efficient and tends to generalize better |
| than point-wise tokenization. |
| |
| 3. A standard transformer encoder processes the patch tokens, and a |
| linear head maps the final representation directly to the |
| forecast horizon. |
| |
| CONTEXT_LEN and HORIZON are fixed (96 and 42) to match the evaluation |
| setup already used for SARIMA/Prophet/statsforecast/t0-alpha in the |
| sibling demand-forecasting project, so results are directly comparable. |
| """ |
| import torch |
| import torch.nn as nn |
|
|
| CONTEXT_LEN = 96 |
| HORIZON = 42 |
| PATCH_SIZE = 8 |
| N_PATCHES = CONTEXT_LEN // PATCH_SIZE |
| D_MODEL = 32 |
| N_HEADS = 4 |
| N_LAYERS = 2 |
| DIM_FEEDFORWARD = 64 |
|
|
|
|
| class TinyForecaster(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.patch_embed = nn.Linear(PATCH_SIZE, D_MODEL) |
| self.pos_embed = nn.Parameter(torch.randn(1, N_PATCHES, D_MODEL) * 0.02) |
|
|
| encoder_layer = nn.TransformerEncoderLayer( |
| d_model=D_MODEL, nhead=N_HEADS, dim_feedforward=DIM_FEEDFORWARD, |
| batch_first=True, dropout=0.1, |
| ) |
| self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=N_LAYERS) |
|
|
| self.head = nn.Sequential( |
| nn.Linear(D_MODEL * N_PATCHES, 128), |
| nn.ReLU(), |
| nn.Linear(128, HORIZON), |
| ) |
|
|
| def forward(self, context_raw: torch.Tensor) -> torch.Tensor: |
| """ |
| context_raw: (batch, CONTEXT_LEN) raw (unnormalized) values. |
| Returns: (batch, HORIZON) forecast, in the ORIGINAL scale. |
| """ |
| mean = context_raw.mean(dim=1, keepdim=True) |
| std = context_raw.std(dim=1, keepdim=True).clamp(min=1e-3) |
| normed = (context_raw - mean) / std |
|
|
| patches = normed.view(-1, N_PATCHES, PATCH_SIZE) |
| tokens = self.patch_embed(patches) + self.pos_embed |
| encoded = self.encoder(tokens) |
| flat = encoded.reshape(encoded.size(0), -1) |
| forecast_normed = self.head(flat) |
|
|
| forecast = forecast_normed * std + mean |
| return forecast |
|
|
