From 1c2ee1d7c313144e4b4ee23b9c08bfbeee4f9249 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 18 May 2026 13:20:26 +0200 Subject: [PATCH] spec(ml-alpha): v2 multi-horizon redesign (A+B+C+D+E integrated) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Integrated design spec for the post-A/B redesign: Kendall sigma BCE (A), L2 anchor on horizon tokens + shared Q (B), horizon-token K-prepend replacing per-horizon Q_h (C), regime-aware MoE gate (D), and inverted-axis attention pass (E). Bundled per pearl_no_deferrals_for_complementary_fixes — all five axes have orthogonal architectural scope and independent kill criteria. Spec drops the C21-C25 per-horizon Q_h path (falsified by sweep 2026-05-18: mean_auc -0.019 vs single-Q baseline) and migrates the existing init buffers into the new horizon-tokens prefix. Co-Authored-By: Claude Opus 4.7 --- ...-05-18-ml-alpha-v2-multi-horizon-design.md | 267 ++++++++++++++++++ 1 file changed, 267 insertions(+) create mode 100644 docs/superpowers/specs/2026-05-18-ml-alpha-v2-multi-horizon-design.md diff --git a/docs/superpowers/specs/2026-05-18-ml-alpha-v2-multi-horizon-design.md b/docs/superpowers/specs/2026-05-18-ml-alpha-v2-multi-horizon-design.md new file mode 100644 index 000000000..b6eb98b3c --- /dev/null +++ b/docs/superpowers/specs/2026-05-18-ml-alpha-v2-multi-horizon-design.md @@ -0,0 +1,267 @@ +# ml-alpha v2 — Multi-Horizon Architecture Redesign + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:writing-plans next to derive the implementation plan from this spec. + +**Goal.** Replace the empirically falsified per-horizon Q_h attention pool (A/B sweep 2026-05-18, mean_auc 0.7559 ± 0.0068 vs baseline 0.7749 ± 0.024) with a single integrated architecture that incorporates five compounding improvements drawn from modern multi-horizon forecasting literature: + +| Axis | Idea | Reference | +|---|---|---| +| **A** | Kendall homoscedastic-uncertainty σ-weighted BCE on the per-horizon losses | [Kendall+Gal+Cipolla 2018](https://arxiv.org/abs/1705.07115), [UW-SO 2025](https://link.springer.com/article/10.1007/s11263-025-02625-x) | +| **B** | L2 anchor regularization on horizon embeddings + the shared Q toward baseline initialization | (signal-driven; floor decays with training progress) | +| **C** | Horizon-token K-prepend — concatenate N_HORIZONS learned tokens to the K-sequence and use a single shared attention, in place of per-horizon Q_h | TFT-style horizon conditioning + standard transformer prefix-tokens | +| **D** | Regime-aware Mixture-of-Experts gate over the pooled context, gated by spread/volatility regime features | [Wavelet-MoE 2025](https://arxiv.org/abs/2508.08825), [Task-aware MoE 2025](https://arxiv.org/html/2509.22279v1) | +| **E** | Inverted attention pass — a second attention block where HIDDEN_DIM features are treated as tokens and time as the embedding dimension | [iTransformer 2024](https://www.datasciencewithmarco.com/blog/itransformer-the-latest-breakthrough-in-time-series-forecasting) | + +These axes are **architecturally independent** (orthogonal scopes) so bundling them is safe per `pearl_no_deferrals_for_complementary_fixes`. Each axis has its own kill-criterion (§Falsifiable claims) so ablations remain meaningful even when run together. + +--- + +## 1. Empirical motivation + +A/B sweep 2026-05-18 with the simple Q_h design (commit `83546b5c3`): + +| Metric | Per-horizon Q_h (n=3) | Baseline single-Q (n=3) | Δ | +|---|---|---|---| +| mean_auc | 0.7559 ± 0.0068 | 0.7749 ± 0.024 | **−0.019** | +| auc_h6000 | 0.7588 ± 0.0049 | 0.7591 ± 0.018 | −0.0003 | +| best_epoch (val_loss) | 1 in 2/3 folds | varies | overfit | +| best_mean_auc_epoch | 4-14 | varies | rank holds | + +Two facts stand out: + +1. `best_epoch=1` on `val_loss` while `best_mean_auc_epoch` is later → val_loss climbs as α opens; ranking quality preserved. **Calibration drift, not rank collapse.** This implicates the loss balancing, not the architecture. + +2. No horizon-distribution shift toward h6000 or any single horizon: spread of per-horizon AUC values is similar to single-Q. **Adding per-horizon queries did not unlock horizon-specific representation** — the existing GRN heads already saturate that axis. + +Modern SOTA architectures (TFT, PatchTST, iTransformer, LiT, TLOB) **do not use per-horizon queries**. The wins in 2024-2025 come from (i) loss balancing, (ii) prefix-token horizon conditioning, (iii) regime-aware routing, (iv) inverted-axis attention as a second pass. + +--- + +## 2. Architecture overview + +``` + MBP-10 snapshots + │ + snap_feature_assemble [F_DIM] + │ + VSN (TFT) [HIDDEN_DIM] + │ + Mamba2 stack-1 [B, K, H] + │ + LayerNorm-a + │ + Mamba2 stack-2 [B, K, H] + │ + LayerNorm-b [B, K, H] = X + │ + ┌───────────────────┴───────────────────┐ + │ │ + (a) standard attention pool (e) inverted attention + with horizon-prefix tokens HIDDEN_DIM features as tokens + (replaces per-horizon Q_h) time as embedding + Out: ctx_h [B, N_H, H] Out: feat_inv [B, H, K] → pool to [B, H] + │ │ + └───────────────┬───────────────────────┘ + │ concat + ▼ + fused_ctx [B, N_H, 2H] + │ + (d) regime-MoE gate routes by regime_features + │ + routed_ctx [B, N_H, H] + │ + GRN heads (existing) + │ + logit [K, B, N_H] + │ + (a) Kendall σ-weighted BCE per horizon + (b) + L2 anchor on horizon tokens + shared Q + │ + total loss +``` + +--- + +## 3. Detailed design + +### 3.1 (C) Horizon-token K-prepend — replaces per-horizon Q_h + +Currently we have `Q_h[N_HORIZONS, HIDDEN_DIM]` and run N_HORIZONS separate attention passes, each with its own Q. + +**v2**: introduce `horizon_tokens[N_HORIZONS, HIDDEN_DIM]` — N_HORIZONS learned tokens that get **prepended** to the K-sequence before attention. A single shared learned query `Q[HIDDEN_DIM]` then attends over `[horizon_tokens; LN_b_out]` of shape `[B, N_HORIZONS + K, HIDDEN_DIM]`. + +Per-horizon outputs come from reading the post-attention state at each horizon-token position: +``` +extended = [horizon_tokens; LN_b_out] # [B, N_H + K, H] +scores = Q · extended # [B, N_H + K] +attn = softmax(scores) # [B, N_H + K] +ctx_h[b, h, :] = attn[b, h] * horizon_tokens[h, :] + + Σ_k attn[b, N_H + k] * LN_b_out[b, k, :] +``` + +This is the prefix-token pattern used in BERT (`[CLS]`) and TFT. Per-horizon conditioning emerges from the token's content, not from a separate Q. Token count: N_HORIZONS × HIDDEN_DIM = 640 params, identical to the old `Q_h` count. + +**Why it's better than per-horizon Q_h:** +- Single shared attention forces the model to learn ONE attention pattern that is then *modulated* by per-horizon tokens. Less capacity, less overfit. +- Tokens can also encode static metadata (horizon length, target type) as learned embedding. +- Removes the per-horizon attention-pool kernel; only one attention kernel needed. + +**Kernel.** A new `horizon_token_attention_pool` kernel that takes the extended key sequence and emits ctx_h. Replaces `per_horizon_attention_pool_*`. + +### 3.2 (E) Inverted attention pass + +In standard attention K positions are tokens and HIDDEN_DIM is the embedding. iTransformer inverts: HIDDEN_DIM features become tokens (variates), and the K-sequence becomes the embedding dimension. Attention then captures cross-variate dependencies that standard attention misses. + +For our LN_b_out `[B, K, H]`: +``` +# Inverted view: each of H features is a token; its embedding is its K-trajectory. +X_inv = LN_b_out.transpose(1, 2) # [B, H, K] +scores = (X_inv @ Q_inv) / sqrt(K) # [B, H, H] cross-feature attention +attn = softmax(scores) # [B, H, H] +feat_inv[b, h, :] = Σ_j attn[b, h, j] * X_inv[b, j, :] # [B, H, K] +# Pool feat_inv across K (mean pool) to recover [B, H]. +``` + +Final inverted output: `pooled_inv [B, H]`, then broadcast to `[B, N_HORIZONS, H]` (same vector per horizon — the inverted pass is horizon-agnostic; it captures regime-shared structure). + +**Kernel.** New `inverted_attention_pool` kernel. Block-per-batch, processes H × H attention. + +**Concat with horizon-token attention** to form `fused_ctx [B, N_HORIZONS, 2H]`, then a `1×1` projection to `[B, N_HORIZONS, H]` before the regime gate. + +### 3.3 (D) Regime-aware MoE gate + +Per memory `pearl_snapshot_alpha_is_regime_conditional`, snapshot-level alpha concentrates in ~20% of book states (spread quintile, volatility regime, time-of-day). A single dense head averages over regimes. + +**Regime features** (computed in `snap_feature_assemble`, broadcast to per-window): +- Spread quintile (1 of 5): one-hot from snapshot-level bid-ask spread vs rolling quintile boundaries +- Volatility regime (1 of 3): low/medium/high from rolling realized vol +- Time bucket (1 of 4): session-open / mid-morning / mid-afternoon / close +Total: 12-dim one-hot regime feature `R [B, 12]`. + +**MoE gate**: +``` +gate_logits = W_gate @ R # [B, N_EXPERTS] (N_EXPERTS = 4) +gate_probs = softmax(gate_logits) # [B, N_EXPERTS] +# Each expert is a small W_e [HIDDEN_DIM, HIDDEN_DIM] linear: +expert_outs[e] = W_e @ fused_ctx # [B, N_HORIZONS, HIDDEN_DIM] +# Sparse top-1 routing (single-expert per batch). +top_e = argmax(gate_probs, axis=-1) # [B] +routed_ctx[b] = expert_outs[top_e[b]] +``` + +Top-1 is preferred over soft (top-K weighted sum) because (i) it scales linearly with N_EXPERTS, (ii) it's reproducible across runs, (iii) it's what the [Wavelet-MoE 2025](https://arxiv.org/abs/2508.08825) paper uses for financial time series. + +**Load-balancing.** Add an auxiliary `aux_loss = N_EXPERTS · Σ_e (frac_e × prob_e_mean)` term per the Switch Transformer recipe to prevent expert collapse. + +**Kernel.** New `regime_moe_gate` kernel; computes gate_logits + softmax + top-1 + expert dispatch. Bwd routes gradients only through the selected expert per batch. + +### 3.4 (A) Kendall σ-weighted BCE + +Replace `total_loss = Σ_h BCE(p_h, y_h)` with: +``` +total_loss = Σ_h [ (1 / (2σ_h²)) · BCE(p_h, y_h) + log σ_h ] +``` +where `σ_h` are learnable per-horizon scalars (parameterised as `log_sigma_h` for stability — Adam updates log_sigma; `σ_h = exp(log_sigma_h)`). + +This auto-balances horizons: a horizon with high BCE (hard) gets large σ_h → its loss is downweighted; the `log σ_h` term prevents σ from running to infinity. + +**Implementation.** Modify `bce_loss_multi_horizon` kernel to read `log_sigma_h [N_HORIZONS]`, compute `(1/(2*exp(2*log_sigma_h))) · bce + log_sigma_h`. Add a new param group with its own AdamW (LR scaled 1/4 of the head LR — same convention as α in the per-horizon design). + +### 3.5 (B) L2 anchor regularization + +Anchor the learnable horizon-token embeddings + shared Q toward their initialization values: +``` +anchor_loss = λ_anchor(t) · ( ||horizon_tokens - horizon_tokens_init||² + ||Q - Q_init||² ) +``` +λ_anchor(t) is signal-driven (no tuned constant): start with `λ_anchor(0)` = anchor_floor (computed from initial parameter magnitudes), decay via a Wiener-α controller as training progresses and the loss improvement plateaus. Floor at `λ_anchor_floor` to prevent total decay (per `pearl_blend_formulas_must_have_permanent_floor`). + +**Init values** are captured at trainer construction in a buffer (per `pearl_first_observation_bootstrap` — sentinel = 0; first observation replaces directly). + +**Kernel.** A small `anchor_l2` kernel computes Σ (p - p_init)² over the flat parameter array; trainer aggregates the per-param-group L2 into `anchor_loss` and adds to total. + +--- + +## 4. Parameter count delta + +| Component | v1 (per-horizon Q_h) | v2 | +|---|---|---| +| Per-horizon attention pool | Q_h: 5×128 = 640 | (removed) | +| Per-horizon residual head | 5×128 + 5 = 645 | (removed) | +| α-gate | 5 | (removed) | +| Horizon tokens | — | 5×128 = 640 | +| Shared Q (already exists) | 128 (kept) | 128 (kept) | +| Inverted-attention Q_inv | — | 32×128 = 4096 (8 cross-feature heads × 128 dim) | +| MoE gate W_gate | — | 12×4 = 48 | +| MoE experts W_e (×4) | — | 4 × 128×128 = 65,536 | +| Fused 1×1 projection W_fuse | — | 2H × H = 256×128 = 32,768 | +| log_sigma per horizon | — | 5 | +| Init buffers (anchor) | — | 640 + 128 = 768 (non-trainable) | +| **Total new trainable** | **1,290** | **~103K** | + +v2 is ~80× more parameters but distributed across regularized and routed paths. Specifically the MoE experts dominate (~65K of 103K) but only one is active per batch via top-1 routing. + +--- + +## 5. Falsifiable claims + +The integrated design wins iff at least 3 of 5 axis-level criteria below pass on the 3-fold A/B vs baseline (task #200): + +| Axis | Falsifiable claim | Threshold | +|---|---|---| +| **A** (Kendall σ) | val_loss best_epoch is no longer epoch 1 (calibration drift fixed) | `mean(best_epoch) ≥ 3` | +| **B** (L2 anchor) | val_loss does not climb past epoch 5 (anchor prevents drift) | `val_loss[e=10] ≤ val_loss[e=2] · 1.10` | +| **C** (horizon tokens) | per-horizon AUC variance increases (horizons differentiate) | `var(best_mean_auc_per_horizon) ≥ 1.5 × baseline var` | +| **D** (regime-MoE) | regime-conditional AUC lift in spread-Q4 | `auc[spread-Q4] − auc[all] ≥ +0.05` | +| **E** (inverted attn) | mean_auc lifts on horizons where cross-variate info matters | `mean_auc h1000+h6000 ≥ baseline + 0.005` | + +Additionally — **wall-time gate**: ≤ 30 s/epoch with the warp-shuffle perf rewrite (i.e., ≤ 2× the 17 s baseline). Anything beyond that triggers further perf work before deploying. + +--- + +## 6. Smoke + cluster validation gates + +1. **Local numgrad parity** for every new kernel (horizon_token_attention_pool, inverted_attention_pool, regime_moe_gate, anchor_l2, modified bce). All within `5e-2` rel-tol / `5e-3` abs-floor. +2. **Local end-to-end smoke** (the existing `per_horizon_full_pipeline_smoke` extended for v2 — α=0 → bit-identical-to-baseline at init must STILL hold). +3. **Cluster smoke** (1 epoch × 1000 train seqs) on the new commit before A/B. +4. **3-fold cluster A/B** at 30 epochs × default 8000 train seqs. + +--- + +## 7. Migration discipline + +Per `feedback_no_partial_refactor.md`, every consumer migrates atomically. The C24/C25 per-horizon path (`PerHorizonAttentionPool`, `PerHorizonResidualHead`, `prob_blend` kernels) gets **fully removed** in the same commit that wires the v2 path. Kernels and source files are deleted, not deprecated. The horizon-tokens repurpose the existing Q_h initialization (already at 1/√HIDDEN_DIM scale per spec C21 §5). + +Compatibility note: cluster smoke at v2 commit will rebuild the binary from scratch; no shared binary cache with v1 (different cubin set). + +--- + +## 8. Out of scope + +- LoRA / micro-LoRA expert adapters (D's experts are full 128×128 — small enough that LoRA buys nothing) +- Quantile loss / pinball loss (binary outcome, not quantile regression) +- LayerNorm on horizon tokens (let them learn their own scale) +- Cross-horizon attention (e.g., horizon h6000 attending to h30 token) — adds N²_H complexity without clear ROI per literature +- Per-expert dropout (handled by MoE's stochastic routing) + +--- + +## 9. Open questions (logged, not blocking) + +1. N_EXPERTS = 4 — sweep candidate later if regime conditionality is stronger than expected +2. Inverted-attention head count (8 vs 4 vs 16) — sweep candidate +3. Anchor floor floor: derive from `||init_param||/100` vs fixed `1e-4` — adopt signal-driven default + +--- + +## 10. References + +- [TFT — Lim+Arik+Loeff+Pfister 2019](https://arxiv.org/abs/1912.09363) +- [PatchTST — Nie+Nguyen 2023](https://arxiv.org/abs/2211.14730) +- [iTransformer — Liu et al. 2024](https://www.datasciencewithmarco.com/blog/itransformer-the-latest-breakthrough-in-time-series-forecasting) +- [LiT — Frontiers 2025](https://www.frontiersin.org/journals/artificial-intelligence/articles/10.3389/frai.2025.1616485/full) +- [TLOB dual-attention 2025](https://github.com/LeonardoBerti00/TLOB) +- [Kendall+Gal+Cipolla 2018](https://arxiv.org/abs/1705.07115) +- [UW-SO 2025](https://link.springer.com/article/10.1007/s11263-025-02625-x) +- [Wavelet-MoE 2025](https://arxiv.org/abs/2508.08825) +- [Task-aware MoE 2025](https://arxiv.org/html/2509.22279v1) +- [Switch Transformer (load balancing) 2021](https://arxiv.org/abs/2101.03961)