foxhunt/docs/superpowers/specs/2026-05-18-per-horizon-attention-pool-design.md

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:

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<f32> } for single-query (legacy).
• pub struct PerHorizonAttentionPool { q_h_d: CudaSlice<f32> } for the new path.
• Both implement a common pub trait AttentionPoolKernel { fn forward(...) -> Result<CudaSlice<f32>>; ... } so the trainer can swap them via a config flag.

crates/ml-alpha/src/trainer/perception.rs config flag:

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<Vec<f32>> ([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