diff --git a/docs/superpowers/specs/2026-05-18-per-horizon-attention-pool-design.md b/docs/superpowers/specs/2026-05-18-per-horizon-attention-pool-design.md new file mode 100644 index 000000000..8db254072 --- /dev/null +++ b/docs/superpowers/specs/2026-05-18-per-horizon-attention-pool-design.md @@ -0,0 +1,204 @@ +# Per-Horizon Attention Pool — Design Spec + +**Status:** DESIGN +**Date:** 2026-05-18 +**Branch (future):** `ml-alpha-phase-c` (separate from `ml-alpha-phase-a` real-LOB work) +**Owner:** ml-alpha team +**Tracked tasks:** #203 (kernels + wiring) + #204 (A/B sweep) + +## 0. Context + +`pearl_phase1d2_decisive_pass` (2026-05-15) validated the Phase 1+2+3 architecture at K=6000 with a **single shared attention query** `Q[HIDDEN_DIM]` pooling LN_b's `[B, K, HIDDEN_DIM]` output into one context vector that initialises the CfC k=0 state. Phase B (`ml-alpha-phase-a`) refactored the kernel to block-per-batch and 3-fold CV at `mean_auc = 0.7749 ± 0.024 / h6000 0.7591 ± 0.018`. The original brainstorm called out per-horizon attention pool variants as a "separate model-side spec". + +**Hypothesis:** different horizons attend to different parts of the K-sequence. A buy-now (h30) signal lives in the most-recent 5-10 snapshots; a 2.4hr (h6000) signal lives in regime-defining state from the entire K window. One shared query has to compromise; per-horizon queries can specialise. + +**Falsifiable claim:** per-horizon attention pool lifts h6000 val AUC by ≥ +0.01 absolute (within the noise floor of single-fold AUC variance per `pearl_single_window_oos_is_not_oos.md`) on a held-out 3-fold CV, OR — at minimum — shifts the per-horizon AUC distribution so short horizons (h30, h100) gain at least as much as h6000 loses. The decision criterion is total cross-fold mean AUC averaged over horizons weighted by Phase-1d.4 deployment economics (currently h6000 = 1.0, others = 0.0; will broaden as cost-frontier sweeps surface other viable horizons). + +Out of scope: +- Multi-head attention (more than per-horizon — separate spec). +- Cross-attention to external context (no-op for single-instrument ES). +- Position embeddings for K (orthogonal — could land after per-horizon). + +## 1. Architecture + +LN_b output: `[B, K, HIDDEN_DIM]` (unchanged). + +**Current** (single query): +``` +Q: [HIDDEN_DIM] shared learned param +scores[b, k] = Σ_h Q[h] * LNb[b, k, h] [B, K] +attn[b, k] = softmax_k(scores[b]) [B, K] +context[b, h] = Σ_k attn[b, k] * LNb[b, k, h] [B, HIDDEN_DIM] +``` +context → CfC k=0 state → K-loop recurrence → multi-horizon heads. + +**Proposed** (per-horizon queries): +``` +Q_h: [N_HORIZONS, HIDDEN_DIM] five learned queries +scores_h[b, h, k] = Σ_d Q_h[h, d] * LNb[b, k, d] [B, N_HORIZONS, K] +attn_h[b, h, k] = softmax_k(scores_h[b, h]) [B, N_HORIZONS, K] +context_h[b, h, d] = Σ_k attn_h[b, h, k] * LNb[b, k, d] [B, N_HORIZONS, HIDDEN_DIM] +``` + +Two consumption paths to pick from in the spec-review: + +### Path A (recommended) — per-horizon contexts feed heads directly, CfC unchanged +- The CfC k=0 state is still initialised by the **mean of per-horizon contexts** (preserves the existing K-loop recurrent semantics without disruption). +- The multi-horizon heads kernel changes to read each horizon's own `context_h[b, h, :]` as the *direct* input (bypassing the CfC's final h_K for the head's primary signal). +- Heads still consume CfC h_K as a residual / auxiliary input (concatenated, then projected back to HIDDEN_DIM) so the K-loop signal doesn't vanish. +- **Pro:** CfC stays load-bearing; the K-loop benefit (state amplification per `pearl_state_amplifies_short_horizon_into_long_horizon`) is preserved. +- **Pro:** Minimal disruption to existing kernels — only attn_pool + multi_horizon_heads change shape; CfC + LN + Mamba2 + VSN unchanged. +- **Con:** 5× attention compute + memory (manageable: K=6000, HIDDEN_DIM=128, FP32 → ~30 MB extra activation memory per batch, well within L40S budget). + +### Path B — separate CfC per horizon +- Each horizon gets its own context AND its own CfC instance with separate weights. +- 5× the CfC compute + 5× the weight count. +- **Pro:** Maximum specialisation. +- **Con:** Capacity explosion (CfC weights are HIDDEN_DIM² = ~16k; 5× = 80k extra) + needs careful regularisation. High train-time cost. + +### Path C — drop CfC, per-horizon contexts feed heads directly +- Skip the K-loop entirely; rely on Mamba2-l2 + LN_b for the temporal modelling. +- **Pro:** Simplest. +- **Con:** Loses the K-loop / CfC benefit empirically validated in Phase 1+2+3. Likely regression. + +**Default for v1: Path A.** Path B is a v2 refinement if Path A doesn't move the needle. + +## 2. Kernel changes + +Two new CUDA kernels replace `attention_pool.cu`: + +`per_horizon_attention_pool_fwd`: +```cuda +grid_dim = (B, N_HORIZONS, 1) +block_dim = (HIDDEN_DIM=128, 1, 1) +``` +- Per-block: one (batch, horizon) pair handles HIDDEN_DIM threads → one context vector. +- Block tree-reduce over HIDDEN_DIM for the per-k dot product (existing pattern from `attention_pool.cu`). +- Per-horizon softmax over K (per-block within the (b, h) slice). +- Final context[b, h, h_dim=tid] accumulation. +- Per `feedback_no_atomicadd.md`: block tree-reduce only. + +`per_horizon_attention_pool_bwd`: +- Per-block (B, N_HORIZONS): single-writer discipline for `d_Q_h` per (h, dim). +- `d_Q_h` is `[N_HORIZONS, HIDDEN_DIM]` shared across all batches. Per-block grad scratch is `[B, N_HORIZONS, HIDDEN_DIM]` then reduced via `reduce_axis0` (existing kernel from Phase B C1). +- `d_LNb[b, k, d]` is per-batch indexed; each (b, h) block adds its contribution via plain `+=` (no race because batches are disjoint, horizons within a batch are summed serially within the block). + +`multi_horizon_heads_fwd` extends signature: +- Old: `(LNb_pooled[B, HIDDEN_DIM]) → probs[B, N_HORIZONS]`. +- New: `(per_horizon_contexts[B, N_HORIZONS, HIDDEN_DIM], cfc_h_K[B, HIDDEN_DIM]) → probs[B, N_HORIZONS]`. +- Each horizon head reads `context[b, h, :]` concat `cfc_h_K[b, :]` (residual) → projects to scalar logit. +- Head weights: `head_w[N_HORIZONS, 2 * HIDDEN_DIM]` (was `[N_HORIZONS, HIDDEN_DIM]`). + +`multi_horizon_heads_bwd` mirrors the above with the existing block-per-batch reduction pattern from Phase B C2. + +## 3. Rust-side changes + +`crates/ml-alpha/src/attention_pool.rs` (new — extract the existing inline binding from cfc/trunk.rs first if needed): +- `pub struct AttentionPool { q_d: CudaSlice }` for single-query (legacy). +- `pub struct PerHorizonAttentionPool { q_h_d: CudaSlice }` for the new path. +- Both implement a common `pub trait AttentionPoolKernel { fn forward(...) -> Result>; ... }` so the trainer can swap them via a config flag. + +`crates/ml-alpha/src/trainer/perception.rs` config flag: +```rust +pub struct PerceptionTrainerConfig { + // ... existing fields ... + pub attention_pool_variant: AttentionPoolVariant, +} + +pub enum AttentionPoolVariant { + SharedQuery, // current — Phase 1+2+3 default + PerHorizonQuery, // new +} +``` + +Default stays `SharedQuery` to preserve all existing checkpoint compatibility. The new variant is opt-in via `--attention-pool per-horizon` CLI flag. + +`CheckpointV1` bumps to `CheckpointV2`: +- Add `attention_pool_variant: u8` discriminant. +- Add `q_h: Option>` ([N_HORIZONS × HIDDEN_DIM] when PerHorizon). +- Load path branches: V1 file → SharedQuery, V2 file → respect the discriminant. + +## 4. Validation strategy + +Three rings, in increasing rigour: + +### Ring 1: kernel-level numgrad parity +- Construct a tiny K=16, B=2, N_HORIZONS=5 fixture. +- Run forward via the new kernels; capture probs. +- Run backward; verify gradients match numerical FD-grad within `1e-4` tolerance (same approach as `Phase 2D.2 VSN backward numgrad check`). + +### Ring 2: smoke training parity +- Branch off the Phase 1+2+3 baseline checkpoint. +- Train one epoch with `--attention-pool per-horizon`; assert no NaN, loss decreases, val AUC > chance (≥ 0.51). +- This is a wiring check, not a quality check. + +### Ring 3: A/B sweep (tracked as task #204) +- 3-fold CV comparison: + - Branch A: `--attention-pool shared-query` (existing baseline) + - Branch B: `--attention-pool per-horizon` +- 30 epochs each, identical seeds + folds + data. +- Compare mean_auc per horizon + total. Per the falsifiable claim in §0, ship Branch B only if it lifts h6000 by ≥ +0.01 OR shifts the per-horizon distribution toward viable horizons (h1000, h300) with no net loss at h6000. + +## 5. Risks + open questions + +| Risk | Severity | Mitigation | +|---|---|---| +| 5× attention compute slows training meaningfully | Medium | Profile early. If wall regression > 20 %, reorder the attention loop to amortise (per-K row read once across all 5 horizons in the same block). | +| Per-horizon Q_h params under-trained at short K | Low | First-observation bootstrap pattern (`pearl_first_observation_bootstrap.md`) doesn't apply to weight params; standard Xavier init + Phase 1d learning rate schedule. | +| Auxiliary CfC h_K input creates spurious correlation that swamps per-horizon signal | Medium | Path A config flag `--head-aux-weight` defaults to 0.25 (heads consume 75% per-horizon + 25% CfC h_K). Tune if needed. | +| Checkpoint compatibility breaks for existing trained models | Low | V1 → V2 migration is read-only (V1 files load as SharedQuery; new training writes V2 with discriminant). | +| Bytecode VM in real-LOB sim doesn't know about per-horizon contexts | Out of scope | Real-LOB consumes only the inference `probs[N_HORIZONS]` output, which has the same shape in both variants. No backtest-side changes needed. | + +### Open questions + +1. **Q_h initialisation**: Xavier scaled to `1/sqrt(HIDDEN_DIM)` (same as existing Q) vs per-horizon-rescaled scaled to `1/sqrt(N_HORIZONS * HIDDEN_DIM)` (acknowledging the 5× param count). Default to the former; revisit if training-time variance is high. +2. **Aux-weight schedule**: should `--head-aux-weight` decay from 1.0 → 0.0 over the first epoch (annealing) or stay constant? Constant for v1. +3. **K-window length**: Path A assumes K=6000 (matches Phase 1+2+3). Per-horizon attention may benefit from shorter K for short horizons. Out of scope for v1 — fixed at K=6000. + +## 6. Compliance with HARD memory rules + +| Rule | How this spec complies | +|---|---| +| `feedback_no_atomicadd.md` | New per-horizon attention kernels use block tree-reduce, same as existing attention_pool.cu. | +| `feedback_no_htod_htoh_only_mapped_pinned.md` | No new host↔device transfer paths. | +| `feedback_no_nvrtc.md` | All kernels pre-compiled to cubins via existing build.rs in `crates/ml-alpha/`. | +| `feedback_no_cpu_test_fallbacks.md` | Ring 1 numgrad parity is GPU-vs-GPU (FD grad on device). | +| `feedback_no_partial_refactor.md` | Variant toggle is a config flag — old SharedQuery path remains fully functional; new PerHorizon path is opt-in. Checkpoint format bumps V1 → V2 with explicit migration. | +| `pearl_state_amplifies_short_horizon_into_long_horizon.md` | Path A preserves the CfC K-loop entirely; per-horizon contexts are an ADDITIVE signal layered on top, not a replacement. | +| `pearl_no_partial_refactor.md` | New kernels + trainer wiring + checkpoint bump land in a single branch (`ml-alpha-phase-c`); no half-shipped state. | + +## 7. Implementation handoff + +This spec is intentionally implementation-deferred. The decision to invest in per-horizon attention pool should be informed by: + +1. **Time + GPU budget** for a 30-epoch × 3-fold A/B (≈ 6 × Phase 1+2+3 training time → ~3-6 hrs on L40S × 5 GPUs). +2. **Whether the per-horizon cost-frontier sweep** (task #202 follow-up via `fxt-backtest sweep`) has identified additional viable horizons beyond h6000. Per `pearl_phase1d4_backtest_cost_edge_frontier.md` the cost edge for h6000 alone is already tight; broadening horizon viability is what makes per-horizon attention valuable. +3. **The 3-fold variance baseline** — the `mean_auc 0.7749 ± 0.024` band means a +0.01 lift is within noise; a meaningful effect needs ≥ +0.024 or a per-horizon distributional shift that's qualitatively visible. + +If those conditions are met, the next step is invoking `superpowers:writing-plans` against this spec to produce a multi-commit implementation plan (estimated 6-8 commits: 2 kernel pairs forward/backward, perception trainer wiring, checkpoint V2 bump, numgrad Ring 1, smoke Ring 2, A/B sweep Ring 3). + +--- + +## Appendix: file inventory (for plan handoff) + +New files: + +``` +crates/ml-alpha/cuda/per_horizon_attention_pool.cu # NEW — fwd + bwd +crates/ml-alpha/src/attention_pool.rs # NEW — variant trait + impls +crates/ml-alpha/tests/per_horizon_attn_numgrad.rs # NEW — Ring 1 fixture +``` + +Modified files: + +``` +crates/ml-alpha/cuda/multi_horizon_heads.cu # extend forward signature to consume + # per-horizon contexts + CfC h_K residual +crates/ml-alpha/src/heads.rs # head_w shape [N_HORIZONS, 2 * HIDDEN_DIM] +crates/ml-alpha/src/cfc/trunk.rs # bump CheckpointV1 → V2 with + # attention_pool_variant + optional q_h +crates/ml-alpha/src/trainer/perception.rs # AttentionPoolVariant config + wiring +crates/ml-alpha/examples/alpha_train.rs # --attention-pool CLI flag +infra/k8s/argo/alpha-perception-template.yaml # --attention-pool plumb-through +scripts/argo-alpha-perception.sh # add --attention-pool flag handling +```