| """Modular-TS training model: a Timestral-style network trained from scratch with DHF. |
| |
| The network *is* `timestral.model.TimestralModel` (Apache-2.0 architecture code, random |
| init), so parameter names match the open inference package and a trained checkpoint |
| exports to `config.json` + `model.safetensors` that `timestral.TimestralForecaster` |
| loads unchanged. We add `forward_train`, the Decoupled-Horizon-Forecasting training pass |
| (our reconstruction; Mistral has not published it): |
| |
| - batch layout (B groups, V rows per group, T = N * patch steps); rows are target, |
| past-only covariate, known-future covariate, or padding; |
| - Enc = all N patch embeddings (causal per patch, past-only scaling + asinh); |
| Dec = N copies of the learned forecast query; per group a min-context s, decoder |
| blocks of K patches anchored at s, s+K, ...; Dec token in block b reads Enc[< b] plus |
| the causal Dec chain of its own block (timestral.masks.build_two_stream_time_mask); |
| - variate attention per patch slot over the V rows of a group with the covariate |
| firewall (covariates never read targets; known-future rows never read causal rows); |
| - target of Dec token p = true patch p normalized in the causal frame of patches < p |
| (the "advancing refine frame" the inference code decodes in), asinh-squashed; |
| - loss = pinball over the 21 quantile knots on observed target values of active patches. |
| |
| NEVER load Timestral weights into this model (Modified-MIT licence, see NOTICE). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| import os |
| import sys |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import Tensor |
|
|
| _HERE = os.path.dirname(os.path.abspath(__file__)) |
| _TS_SRC = os.path.join(os.path.dirname(_HERE), "timestral", "src") |
| if os.path.isdir(_TS_SRC) and _TS_SRC not in sys.path: |
| sys.path.insert(0, _TS_SRC) |
|
|
| from timestral.config import TimestralConfig |
| from timestral.layers import Attention, _apply_rotary |
| from timestral.masks import build_two_stream_time_mask, compute_enc_key_validity |
| from timestral.model import ROLE_KNOWN_COV, ROLE_PAST_COV, ROLE_TARGET, TimestralModel |
| from timestral.scaler import apply_affine_asinh_masked, causal_patch_loc_scale |
|
|
| ROLE_PAD = -1 |
| QUANTILES_21 = (0.01, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, |
| 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 0.99) |
|
|
|
|
| def make_config(size: str = "2.5M", context_length: int = 4096) -> TimestralConfig: |
| presets = { |
| |
| "2.5M": (128, 2, 320, 2), |
| "5M": (192, 3, 512, 2), |
| "28M": (384, 6, 1024, 3), |
| "tiny": (64, 1, 128, 2), |
| } |
| d, h, mlp, k = presets[size] |
| return TimestralConfig( |
| patch_size=32, d_model=d, num_heads=h, mlp_hidden_dim=mlp, |
| block_schedule=("time", "time", "time", "time_variate") * k, |
| quantile_knots=QUANTILES_21, block_size=32, context_length=context_length, |
| rope_theta=1e4, rope_partial_factor=0.5, norm_eps=1e-6, attn_bias=False, mlp_bias=False, |
| time_attn_gate="sdpa_sigmoid", variate_attn_gate="sdpa_sigmoid", minimum_scale=1e-6, |
| use_asinh=True, stabilize_with_global=False, covariate_enabled=True, |
| split_context_horizon_attn=True, attend_missing_patches=True, quantile_width_factor=3.0, |
| ) |
|
|
|
|
| def count_params(model: torch.nn.Module) -> dict: |
| tot = sum(p.numel() for p in model.parameters()) |
| emb = sum(p.numel() for n, p in model.named_parameters() if n.startswith("representation")) |
| head = sum(p.numel() for n, p in model.named_parameters() if n.startswith("head")) |
| return {"total": tot, "embed": emb, "trunk": tot - emb - head, "head": head} |
|
|
|
|
| def _variate_attention(attn: Attention, x: Tensor, mask: Tensor, B: int, V: int) -> Tensor: |
| """Variate attention per (group, slot) over V rows. x: (B*V, S, D); mask: (B*S, 1, V, V).""" |
| R, S, D = x.shape |
| inp = x.view(B, V, S, D).permute(0, 2, 1, 3).reshape(B * S, V, D) |
| q, k, v = attn._qkv(inp) |
| out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_mask=mask) |
| out = out.transpose(1, 2).reshape(B * S, V, D) |
| if attn._gate is not None: |
| out = out * attn._gate(inp).sigmoid() |
| out = attn.wO(out) |
| return out.view(B, S, V, D).permute(0, 2, 1, 3).reshape(R, S, D) |
|
|
|
|
| def _time_attention(attn: Attention, x: Tensor, mask: Tensor, cos: Tensor, sin: Tensor) -> Tensor: |
| q, k, v = attn._qkv(x) |
| q, k = _apply_rotary(q, k, cos, sin) |
| out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_mask=mask) |
| out = out.transpose(1, 2).reshape(x.shape[0], x.shape[1], -1) |
| if attn._gate is not None: |
| out = out * attn._gate(x).sigmoid() |
| return attn.wO(out) |
|
|
|
|
| class ModularTS(TimestralModel): |
| """TimestralModel + DHF training pass (identical parameter names).""" |
|
|
| def __init__(self, config: TimestralConfig) -> None: |
| super().__init__(config) |
| self.register_buffer("q_levels", torch.tensor(config.quantile_knots, dtype=torch.float32), persistent=False) |
| self._init_weights() |
|
|
| def _init_weights(self) -> None: |
| for name, p in self.named_parameters(): |
| if p.dim() == 2: |
| torch.nn.init.normal_(p, std=0.02) |
| if name.endswith("wO.weight") or name.endswith("w3.weight"): |
| p.data.mul_(1.0 / math.sqrt(2 * len(self.config.block_schedule))) |
| elif name.endswith(".bias"): |
| torch.nn.init.zeros_(p) |
| torch.nn.init.normal_(self.representation.forecast_query, std=0.02) |
|
|
| |
| def build_masks(self, observed: Tensor, roles: Tensor, s: Tensor): |
| """observed (B,V,T) bool, roles (B,V) long, s (B,) long -> masks/positions/key validity.""" |
| c = self.config |
| B, V, T = observed.shape |
| P, N, R = c.patch_size, T // c.patch_size, B * V |
| dev = observed.device |
| obs_pp = observed.view(R, N, P).any(-1) |
| not_pad = (roles != ROLE_PAD).view(R) |
| key_valid = compute_enc_key_validity(obs_pp, c.attend_missing_patches) & not_pad[:, None] |
| known = (roles == ROLE_KNOWN_COV).view(R) |
| target = (roles == ROLE_TARGET).view(R) |
| s_rows = s.repeat_interleave(V) |
| time_mask = build_two_stream_time_mask(s_rows, N, c.block_size, key_valid, bidirectional_rows=known) |
| |
| slot_valid = torch.cat([key_valid, not_pad[:, None].expand(R, N)], dim=1).view(B, V, 2 * N) |
| keyok = slot_valid.permute(0, 2, 1)[:, :, None, :] |
| tg = target.view(B, V) |
| kn = known.view(B, V) |
| forbid = ((~tg)[:, :, None] & tg[:, None, :]) | (kn[:, :, None] & (~kn)[:, None, :]) |
| vmask = keyok & ~forbid[:, None] |
| eye = torch.eye(V, dtype=torch.bool, device=dev) |
| vmask = (vmask | eye).reshape(B * 2 * N, 1, V, V) |
| |
| grp_obs = obs_pp.view(B, V, N).any(1) |
| first = torch.where(grp_obs.any(-1), grp_obs.float().argmax(-1), torch.zeros(B, device=dev)).long() |
| p = torch.arange(N, device=dev) |
| pos = (p[None] - first[:, None]).repeat_interleave(V, 0) |
| positions = torch.cat([pos, pos], dim=1) |
| return time_mask, vmask, positions, obs_pp, target, known |
|
|
| |
| def forward_train(self, values: Tensor, observed: Tensor, roles: Tensor, s: Tensor, return_preds: bool = False, |
| teacher: Tensor | None = None, teacher_obs: Tensor | None = None, gt_row: Tensor | None = None, |
| lam: float = 0.5, teacher_q_idx: tuple = (2, 4, 6, 8, 10, 12, 14, 16, 18)): |
| """values (B,V,T) float (anything where unobserved), observed (B,V,T) bool, |
| roles (B,V) long in {-1,0,1,2}, s (B,) min-context patches. Returns (loss, stats).""" |
| c = self.config |
| B, V, T = values.shape |
| P, N, R = c.patch_size, T // c.patch_size, B * V |
| x = torch.where(observed, values, torch.zeros_like(values)).float().view(R, T) |
| obs = observed.view(R, T) |
|
|
| with torch.no_grad(): |
| time_mask, vmask, positions, obs_pp, target_rows, known_rows = self.build_masks(observed, roles, s) |
| loc, scale = causal_patch_loc_scale(x[:, None], obs[:, None], P, c.minimum_scale, False) |
| loc, scale = loc[:, 0], scale[:, 0] |
| scaled = apply_affine_asinh_masked(x, loc, scale, obs, use_asinh=c.use_asinh) |
| scaled, _, _ = self._known_cov_scaling(x, obs, scaled, loc, scale, known_rows) |
| |
| loc_pp = loc.view(R, N, P)[..., 0] |
| sc_pp = scale.view(R, N, P)[..., 0] |
| f_loc = torch.cat([loc_pp[:, :1], loc_pp[:, :-1]], dim=1) |
| f_sc = torch.cat([sc_pp[:, :1], sc_pp[:, :-1]], dim=1) |
| y = torch.asinh((x.view(R, N, P) - f_loc[..., None]) / f_sc[..., None].clamp_min(1e-6)) |
| cnt_before = torch.cumsum(obs.view(R, N, P).sum(-1), 1) - obs.view(R, N, P).sum(-1) |
| p_idx = torch.arange(N, device=x.device) |
| active = p_idx[None] >= s.repeat_interleave(V)[:, None] |
| w = (obs.view(R, N, P) & target_rows[:, None, None] & active[..., None] & (cnt_before >= 2)[..., None]).float() |
| if gt_row is not None: |
| w = w * gt_row.view(R, 1, 1).float() |
| y = torch.where(w > 0, y, torch.zeros_like(y)) |
| if teacher is not None: |
| |
| Qt = teacher.shape[-1] |
| tv = torch.where(teacher_obs.view(R, T, 1), teacher.view(R, T, Qt).float(), torch.zeros(1, device=x.device)) |
| yt = torch.asinh((tv.view(R, N, P, Qt) - f_loc[..., None, None]) / f_sc[..., None, None].clamp_min(1e-6)) |
| wt = (teacher_obs.view(R, N, P) & target_rows[:, None, None] & active[..., None] & (cnt_before >= 2)[..., None]).float() |
| yt = torch.where(wt[..., None] > 0, yt, torch.zeros_like(yt)) |
|
|
| x_enc = self.representation(scaled, obs) |
| x_dec = self.representation.forecast_tokens(R, N) |
| h = torch.cat([x_enc, x_dec.to(x_enc.dtype)], dim=1) |
| cos, sin = self.transformer.rope(h, positions) |
| for layer in self.transformer.layers: |
| for norm, attn in zip(layer.attn_norms, layer.attns): |
| if attn.is_time: |
| h = h + _time_attention(attn, norm(h), time_mask, cos, sin) |
| else: |
| h = h + _variate_attention(attn, norm(h), vmask, B, V) |
| h = h + layer.mlp(layer.mlp_norm(h)) |
| h = self.transformer.out_norm(h) |
| preds = self.head(h[:, N:]).view(R, N, P, -1).float() |
|
|
| if return_preds: |
| return preds, f_loc, f_sc |
| err = y[..., None] - preds |
| tau = self.q_levels.view(1, 1, 1, -1) |
| pin = torch.maximum(tau * err, (tau - 1) * err).mean(-1) |
| denom = w.sum().clamp_min(1.0) |
| loss = (pin * w).sum() / denom |
| with torch.no_grad(): |
| mae = ((y - preds[..., c.median_index]).abs() * w).sum() / denom |
| stats = {"pinball": loss.detach(), "mae_med": mae, "n_tgt": denom} |
| if teacher is not None: |
| pq = preds[..., list(teacher_q_idx)] |
| tq = self.q_levels[list(teacher_q_idx)].view(1, 1, 1, -1) |
| et = yt - pq |
| kd = torch.maximum(tq * et, (tq - 1) * et).mean(-1) |
| kd_loss = (kd * wt).sum() / wt.sum().clamp_min(1.0) |
| loss = (1 - lam) * loss + lam * kd_loss |
| stats["kd"] = kd_loss.detach() |
| return loss, stats |
|
|
