gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

spec(ml-alpha): v2 multi-horizon redesign (A+B+C+D+E integrated)

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-18 13:20:26 +02:00
parent 286ea26e2a
commit 1c2ee1d7c3
1 changed files with 267 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

267
docs/superpowers/specs/2026-05-18-ml-alpha-v2-multi-horizon-design.md Normal file
View File

@@ -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)