# 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<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:
```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<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
```
