From 010445b5df8038b5e87285bb1246c840ce058602 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 21:41:02 +0200 Subject: [PATCH] =?UTF-8?q?docs(plans):=20pivot=20Phase=201.7=20to=20TFT?= =?UTF-8?q?=20GRN=20heads;=20add=20Phase=202C=20(TGN=20=CE=94t=20Fourier)?= =?UTF-8?q?=20+=202D=20(TFT=20VSN)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit User directive 2026-05-17: borrow TFT GRN over the planned 2-layer MLP heads. GRN structure: 2-layer GELU MLP body (eta_2 → eta_1) + GLU gate + main + skip-projection from trunk → final sigmoid. Gives per-horizon "linear vs deeper-transform" gating, matches the regime-conditional alpha pattern (pearl_snapshot_alpha_is_regime_conditional). 5x parameter count vs the 2-layer MLP but the gated residual is exactly what TFT empirically wins on. Phase 2C (TGN Δt Fourier features): 8 sin/cos features of Δt at log-spaced periods [60s, 6s, 600ms, 60ms] appended to snap_features. Critical with decision-stride>1 where Δt varies across positions. Bumps FEATURE_DIM 32→40. Phase 2D (TFT VSN): per-feature softmax-normalised gates at the trunk entry, replacing raw concat of snap_features. Learns to down-weight noisy features per regime (canonical: trade-flow in low-volume, OFI in spread-Q4). 2 new param tensors, 1 new cuda kernel (fwd+bwd). Existing 2-layer MLP kernels from Tasks 1.3/1.4 stay in the cubin as ablation baseline; wired path becomes GRN. Phase 1.7 plan now spells out the full GRN forward + backward chain rule (skip + sigmoid(gate) * main → outer sigmoid), kernel signatures, parameter Xavier init, AdamW × 10 setup, and an extended numerical-grad check covering all 10 new param tensors. Co-Authored-By: Claude Opus 4.7 --- .../2026-05-17-model-capacity-scaleup.md | 294 +++++++++++++++++- 1 file changed, 293 insertions(+), 1 deletion(-) diff --git a/docs/superpowers/plans/2026-05-17-model-capacity-scaleup.md b/docs/superpowers/plans/2026-05-17-model-capacity-scaleup.md index 468ba41cc..feb5148c7 100644 --- a/docs/superpowers/plans/2026-05-17-model-capacity-scaleup.md +++ b/docs/superpowers/plans/2026-05-17-model-capacity-scaleup.md @@ -846,7 +846,238 @@ git add crates/ml-alpha/src/trainer/perception.rs git commit -m "feat(ml-alpha): wire LayerNorm between Mamba2 trunk and CfC K-loop" ``` -### Task 1.7: Replace single-layer heads with 2-layer in PerceptionTrainer +### Task 1.7: Replace single-linear heads with TFT GRN (Gated Residual Network) heads + +**Rationale (2026-05-17 amendment):** The original plan called for 2-layer MLP heads. User directive on the same day pivots to a **TFT Gated Residual Network** (Lim et al. 2021 §3.3): a 2-layer ELU/GELU MLP body wrapped in a sigmoid-gated residual with an explicit skip-projection. The gated residual lets each per-horizon head learn `out = skip + sigmoid(gate) * main`, so the model can fall back to the linear-projection of the trunk in low-edge regimes and switch to the deeper-transform path only where it pays off (canonical motivation: `pearl_snapshot_alpha_is_regime_conditional` — alpha lives in ~20% of book states). + +The existing 2-layer MLP kernels from Tasks 1.3/1.4 stay in the cubin as an **ablation baseline**; the wired path is the new GRN kernel. + +**GRN math (per-horizon k, per-sample b):** + +``` +eta_2[k, m] = ELU(W2[k, m, i] @ h[b, i] + b2[k, m]) # [HIDDEN] → [HEAD_MID] +eta_1[k, m] = W1[k, m, n] @ eta_2[k, n] + b1[k, m] # [HEAD_MID] → [HEAD_MID], linear +gate_lin[k] = W_gate[k, m] @ eta_1[k, m] + b_gate[k] # [HEAD_MID] → scalar +main[k] = W_main[k, m] @ eta_1[k, m] + b_main[k] # [HEAD_MID] → scalar +skip[k] = W_skip[k, i] @ h[b, i] + b_skip[k] # [HIDDEN] → scalar +logit[k] = skip[k] + sigmoid(gate_lin[k]) * main[k] +p[k] = sigmoid(logit[k]) +``` + +Parameter count per horizon: W2 (HIDDEN*HEAD_MID) + b2 (HEAD_MID) + W1 (HEAD_MID*HEAD_MID) + b1 (HEAD_MID) + W_gate (HEAD_MID) + b_gate (1) + W_main (HEAD_MID) + b_main (1) + W_skip (HIDDEN) + b_skip (1) ≈ **HIDDEN*(HEAD_MID+1) + HEAD_MID*(HEAD_MID+3) + 3** = 128*65 + 64*67 + 3 = **12,611 params per horizon × 5 horizons = 63,055** (vs the old single-linear's 5*(128+1) = 645, and the planned 2-layer MLP's 5*(128*64+64+64+1) = 41,925). + +Within the per-horizon path GELU vs ELU: keep GELU since the existing `gelu_act`/`gelu_prime` device helpers are battle-tested via Task 1.5's numerical gradient check; ELU's only practical advantage over GELU here is computational simplicity, which doesn't matter at 64-wide intermediate. + +**Files:** +- Modify: `crates/ml-alpha/cuda/multi_horizon_heads.cu` — add `multi_horizon_heads_grn_fwd_batched` + `multi_horizon_heads_grn_bwd_batched` kernels +- Modify: `crates/ml-alpha/src/heads.rs` — add `HEAD_MID_DIM` constant +- Modify: `crates/ml-alpha/src/trainer/perception.rs` — swap head params + kernel calls +- Modify: `crates/ml-alpha/tests/perception_overfit.rs` — keep existing smokes (must still converge with deeper heads) +- Modify: `crates/ml-alpha/tests/heads_2layer_numgrad.rs` — extend with GRN numerical-gradient check + +- [ ] **Step 1: Add GRN forward kernel** + +In `crates/ml-alpha/cuda/multi_horizon_heads.cu`, append after the 2-layer MLP backward kernel: + +```cuda +// Per-horizon TFT GRN forward. +// Layout matches the 2-layer MLP variant for eta_2 + eta_1, plus: +// z2_out[B, 5, HEAD_MID] — saved eta_1[k, m] (input to gate/main/used for grad_eta_1) +// gate_logit_out[B, 5] — saved gate pre-sigmoid (for gate'(x) = σ(x)(1-σ(x)) in bwd) +// main_out[B, 5] — saved main scalar +// logit_out[B, 5] — saved pre-sigmoid final logit (for p = σ(logit)) +extern "C" __global__ void multi_horizon_heads_grn_fwd_batched( + const float* __restrict__ w1, // [5, HEAD_MID, HIDDEN] + const float* __restrict__ b1, // [5, HEAD_MID] + const float* __restrict__ w2, // [5, HEAD_MID, HEAD_MID] + const float* __restrict__ b2, // [5, HEAD_MID] + const float* __restrict__ w_gate, // [5, HEAD_MID] + const float* __restrict__ b_gate, // [5] + const float* __restrict__ w_main, // [5, HEAD_MID] + const float* __restrict__ b_main, // [5] + const float* __restrict__ w_skip, // [5, HIDDEN] + const float* __restrict__ b_skip, // [5] + const float* __restrict__ h, // [B, HIDDEN] + int n_batch, + float* __restrict__ probs, // [B, 5] + float* __restrict__ z1_out, // [B, 5, HEAD_MID] — pre-GELU eta_2_z + float* __restrict__ a1_out, // [B, 5, HEAD_MID] — post-GELU eta_2 + float* __restrict__ z2_out, // [B, 5, HEAD_MID] — eta_1 + float* __restrict__ gate_logit_out, // [B, 5] + float* __restrict__ main_out, // [B, 5] + float* __restrict__ logit_out // [B, 5] — pre-sigmoid final logit +); +``` + +Implementation: per-block per-sample (block.x = b_idx), threads tile HEAD_MID. Pass 1 = 2-layer MLP body (identical to existing kernel up to eta_1). Pass 2 = compute gate_lin / main / skip into shared mem (using 5 threads + first 5 threads). Pass 3 = thread tid<5 writes `logit = skip + sigmoid(gate_lin) * main`, then `probs = sigmoid(logit)`. + +- [ ] **Step 2: Add GRN backward kernel** + +Chain rule (for clarity; per-horizon k): + +``` +d_logit = grad_probs * p * (1 - p) # sigmoid back +d_skip = d_logit +d_main = d_logit * sigmoid(gate_lin) +d_gate_a = d_logit * main # grad through sigmoid(gate_lin) +d_gate_lin = d_gate_a * sigmoid(gate_lin) * (1 - sigmoid(gate_lin)) + +grad_w_skip[k, i] += d_skip * h[i] +grad_b_skip[k] += d_skip +grad_w_main[k, m] += d_main * eta_1[k, m] +grad_b_main[k] += d_main +grad_w_gate[k, m] += d_gate_lin * eta_1[k, m] +grad_b_gate[k] += d_gate_lin + +d_eta_1[k, m] = d_main * w_main[k, m] + d_gate_lin * w_gate[k, m] + +# eta_1 = W1 @ eta_2 + b1 — pure linear: +grad_w1[k, m, n] += d_eta_1[k, m] * eta_2[k, n] +grad_b1[k, m] += d_eta_1[k, m] +d_eta_2[k, n] = sum_m d_eta_1[k, m] * w1[k, m, n] + +# eta_2 = GELU(z1): +d_z1[k, m] = d_eta_2[k, m] * gelu_prime(z1[k, m]) + +grad_w2[k, m, i] += d_z1[k, m] * h[i] +grad_b2[k, m] += d_z1[k, m] + +# Trunk gradient (ISV lambda-scaled, summed across all k, m): +grad_h[bi, i] = sum_k lambda[k] * ( + d_skip[k] * w_skip[k, i] + + sum_m d_z1[k, m] * w2[k, m, i] +) + carry[i] +``` + +Block layout: HEAD_MID=64 threads per block, ONE block per launch (loops over n_batch internally — same single-writer discipline as 2-layer bwd). + +Signature: + +```cuda +extern "C" __global__ void multi_horizon_heads_grn_bwd_batched( + const float* __restrict__ w1, // [5, HEAD_MID, HEAD_MID] + const float* __restrict__ w2, // [5, HEAD_MID, HIDDEN] -- note: matches the rename below + const float* __restrict__ w_gate, // [5, HEAD_MID] + const float* __restrict__ w_main, // [5, HEAD_MID] + const float* __restrict__ w_skip, // [5, HIDDEN] + const float* __restrict__ probs, // [B, 5] + const float* __restrict__ grad_probs, // [B, 5] + const float* __restrict__ z1, // [B, 5, HEAD_MID] + const float* __restrict__ a1, // [B, 5, HEAD_MID] + const float* __restrict__ z2, // [B, 5, HEAD_MID] + const float* __restrict__ gate_logit, // [B, 5] + const float* __restrict__ main_val, // [B, 5] + const float* __restrict__ logit, // [B, 5] + const float* __restrict__ h, // [B, HIDDEN] + const float* __restrict__ grad_h_carry, // [B, HIDDEN] (nullptr OK) + const float* __restrict__ lambda, // [5] + int n_batch, + float* __restrict__ grad_w1, + float* __restrict__ grad_b1, + float* __restrict__ grad_w2, + float* __restrict__ grad_b2, + float* __restrict__ grad_w_gate, + float* __restrict__ grad_b_gate, + float* __restrict__ grad_w_main, + float* __restrict__ grad_b_main, + float* __restrict__ grad_w_skip, + float* __restrict__ grad_b_skip, + float* __restrict__ grad_h +); +``` + +**IMPORTANT naming nit:** in the 2-layer MLP kernel the "input → mid" matrix is `w1` and "mid → output" is `w2`. In the TFT GRN spec above, the convention flips: `W2` is "input → mid" (first linear, ELU-activated) and `W1` is "mid → mid" (the eta_1 linear). Keep the **MLP convention** (`w1` = first, `w2` = second) for kernel naming so the existing 2-layer kernels carry forward without rename. This means the GRN kernel's `w1` is the FIRST linear (HIDDEN→HEAD_MID, GELU-activated), `w2` is the SECOND linear (HEAD_MID→HEAD_MID, no activation). + +Pseudocode block layout: same passes structure as the 2-layer bwd plus three additional passes for skip / gate / main param grads. Reuse `gelu_prime`, `s_d_z1[]`, `s_d_z2[]` shared-mem patterns. + +- [ ] **Step 3: Register GRN kernels in build.rs** (no change — `multi_horizon_heads` already compiled; appending kernels to the .cu file picks them up on next cache-bust). + +Bump cache-bust comment to v7: + +```rust +// Cache bust v7 (2026-05-17): GRN heads added to multi_horizon_heads.cu. +``` + +- [ ] **Step 4: Add `HEAD_MID_DIM` constant** + +In `crates/ml-alpha/src/heads.rs`: + +```rust +pub const HEAD_MID_DIM: usize = 64; +``` + +- [ ] **Step 5: Replace head params in PerceptionTrainer** + +Delete `heads_w_d`, `heads_b_d`, `opt_heads_w`, `opt_heads_b`, `grad_heads_w_d`, `grad_heads_b_d`. Add (all `CudaSlice` unless noted): + +```rust +heads_w1_d, heads_b1_d, heads_w2_d, heads_b2_d, +heads_w_gate_d, heads_b_gate_d, heads_w_main_d, heads_b_main_d, +heads_w_skip_d, heads_b_skip_d, +grad_heads_w1_d, grad_heads_b1_d, grad_heads_w2_d, grad_heads_b2_d, +grad_heads_w_gate_d, grad_heads_b_gate_d, grad_heads_w_main_d, grad_heads_b_main_d, +grad_heads_w_skip_d, grad_heads_b_skip_d, +opt_heads_w1, opt_heads_b1, opt_heads_w2, opt_heads_b2, +opt_heads_w_gate, opt_heads_b_gate, opt_heads_w_main, opt_heads_b_main, +opt_heads_w_skip, opt_heads_b_skip, // AdamW × 10 +// Forward intermediates saved for backward (per K position, per batch): +z1_per_k_d, a1_per_k_d, z2_per_k_d, // [K, B, 5, HEAD_MID] +gate_logit_per_k_d, main_per_k_d, logit_per_k_d, // [K, B, 5] +// Kernel handles +heads_grn_fwd_fn, heads_grn_bwd_fn, +``` + +- [ ] **Step 6: Xavier init for the new params** + +Per-param fan_in: +- `w1`: HIDDEN → HEAD_MID per (k, m); fan_in = HIDDEN +- `w2`: HEAD_MID → HEAD_MID per (k, m); fan_in = HEAD_MID +- `w_gate`, `w_main`: HEAD_MID → 1 per k; fan_in = HEAD_MID +- `w_skip`: HIDDEN → 1 per k; fan_in = HIDDEN +- All biases: 0.0 + +Use `ChaCha8Rng::seed_from_u64(cfg.seed)` to mirror the existing init pattern. + +- [ ] **Step 7: Wire GRN fwd in K-loop** + +Replace the existing `heads_batched_fn` launch with `heads_grn_fwd_fn`. Pass per-K slot pointers for all 6 new intermediate tensors plus probs. + +- [ ] **Step 8: Wire GRN bwd in K-loop** + +Replace the existing `heads_bwd_batched_fn` launch with `heads_grn_bwd_fn`. The new call takes ISV `lambda[5]` (same pointer as the 2-layer bwd would have used). + +- [ ] **Step 9: Replace zero-grad memsets and AdamW updates** + +10 new accumulators to zero at step start (loop over Vec<&mut CudaSlice>). 10 new AdamW.step() calls (after Mamba2 AdamW). + +- [ ] **Step 10: Eval path** + +`evaluate_batched` re-uses the K-loop forward — the GRN fwd kernel runs in eval too, just without backward. + +- [ ] **Step 11: Smoke test** + +```bash +CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo test -p ml-alpha --test perception_overfit -- --nocapture 2>&1 | tail -20 +``` + +Expected: synthetic overfit converges (loss < 0.001 after 250 steps). If not, suspect chain-rule bug in skip/gate paths — numerical gradient check (Task 1.5 → 1.5b) catches this. + +- [ ] **Step 11b: Extend numerical gradient check** + +In `crates/ml-alpha/tests/heads_2layer_numgrad.rs`, add a `grn_numgrad` test that finite-differences each of the 10 GRN parameter tensors. Relax tolerance to `1e-2` for GRN (deeper chain rule, more accumulated FP error). Use ε=1e-3, central diff. + +- [ ] **Step 12: Commit** + +```bash +git add crates/ml-alpha/cuda/multi_horizon_heads.cu \ + crates/ml-alpha/build.rs \ + crates/ml-alpha/src/heads.rs \ + crates/ml-alpha/src/trainer/perception.rs \ + crates/ml-alpha/tests/heads_2layer_numgrad.rs +git commit -m "feat(ml-alpha): GRN heads (gated residual + skip-projection) replace single-linear (Phase 1.7)" +``` + +### Task 1.7 [DEPRECATED]: Replace single-layer heads with 2-layer in PerceptionTrainer **Files:** - Modify: `crates/ml-alpha/src/heads.rs` (add `HEAD_MID_DIM` constant) @@ -1484,6 +1715,67 @@ Acceptance: fold-1 with stride=4 ≥ Phase 1 fold-1 + 1.5pt h6000 → ship 3-fol After Phase 2A and 2B land, run a fresh 3-fold CV with `--decision-stride 4 --seq-len 64 --mamba2-state-dim 32 --epochs 30 --n-train-seqs 24000`. Compare to Phase 1 3-fold baseline. +### Phase 2C: TGN Δt Fourier Features + +**Phase goal:** Encode the actual elapsed time between consecutive snapshots/positions as Fourier features in the snap-feature stream. **Critical with decision-stride** (Phase 2A): the gap between K-positions becomes irregular, and without Δt encoding the model has no way to distinguish `Δt=2ms` from `Δt=2s`. Borrows from TGN (Rossi et al. 2020): time encoding `φ(Δt) = [cos(ω_k · Δt), sin(ω_k · Δt)]_k` with log-spaced frequencies. Expected lift: +0.5-1.5pt on h6000 in regime-shift folds. ~1-2hr. + +**File map:** + +| File | Action | Responsibility | +|------|--------|----------------| +| `crates/ml-alpha/cuda/snap_feature_assemble.cu` | **MODIFY** | Append 8 Fourier-encoded `Δt` features at the end of the feature vector. `Δt = ts_ns - prev_ts_ns` (already a loader output). Frequencies: 4 octaves log-spaced over `[1ms, 60s]` → `(cos, sin)` for each → 8 features. | +| `crates/ml-alpha/src/cfc/snap_features.rs` | **MODIFY** | Bump `FEATURE_DIM` from 32 → 40. Document the new dim layout. | +| `crates/ml-alpha/src/mamba2_block.rs` | **NO change** — `FEATURE_DIM` is consumed via constant; Mamba2's `in_dim = FEATURE_DIM` flows through automatically. | +| `crates/ml-alpha/tests/snap_feature_assemble.rs` | **MODIFY** | Add a test that `Δt = 1ms` and `Δt = 1s` produce visibly-different (>0.1 in L2) Fourier feature vectors. | + +**Frequency schedule:** `ω_k = 2π / period_k` where periods are `[60s, 6s, 600ms, 60ms]`. This gives sin/cos features sensitive to scales from sub-second microstructure up to minute-scale macro context. + +**Tasks (bite-sized):** + +- Task 2C.1: Add Fourier encoder to `snap_feature_assemble.cu` (8 new feature dims at end of vector) +- Task 2C.2: Bump `FEATURE_DIM` constant, add scale-discrimination test +- Task 2C.3: Re-run fold-1 smoke with Δt features — expect ≥ Phase 2B baseline (no regression, possible +0.5pt at long horizons) + +**Acceptance criteria:** +- Distinct `Δt` values produce L2-distinct Fourier features (test invariant). +- Phase 2C fold-1 ≥ Phase 2B fold-1. Lift not guaranteed at stride=1 (Δt ≈ constant); should appear at stride=4 where Δt has more variation. + +### Phase 2D: TFT Variable Selection Network (VSN) + +**Phase goal:** Replace the raw concat of `snap_features` (32-dim, soon 40-dim with Δt) at the trunk entry with a TFT-style VSN: per-feature gates + softmax normalisation → learned weighted mixture. Lets the model down-weight noisy features in regimes where they aren't informative (canonical: trade-flow features in low-volume regimes; OFI features in spread-Q4-dominant regimes). Borrows from TFT (Lim et al. 2021 §4.2). Expected lift: +1-2pt mean AUC by reducing input noise. ~2-3hr. + +**File map:** + +| File | Action | Responsibility | +|------|--------|----------------| +| `crates/ml-alpha/cuda/variable_selection.cu` | **CREATE** | Per-feature gate kernel: 32 (or 40) input features → 32 (or 40) gate logits → softmax over features → gates ∈ [0,1] summing to 1 → output = gates · features. Forward + backward. | +| `crates/ml-alpha/build.rs` | **MODIFY** | Add `"variable_selection"` to KERNELS, bump cache-bust to v9 (post-Phase-1.7 GRN). | +| `crates/ml-alpha/src/trainer/perception.rs` | **MODIFY** | Apply VSN between `snap_feature_assemble` output and Mamba2 input. New params: `vsn_w [FEATURE_DIM, FEATURE_DIM]` (gate weights), `vsn_b [FEATURE_DIM]` (gate biases). 2 new AdamW. | +| `crates/ml-alpha/tests/perception_overfit.rs` | **MODIFY** | Smoke still uses VSN unconditionally — must converge with VSN gating active (proves the gradient flow is correct). | + +**VSN math:** for each sample b, position k: + +``` +x[b, k, :] = snap_features[b, k, :] # [FEATURE_DIM] +gate_logit[i] = sum_j vsn_w[i, j] * x[b, k, j] + vsn_b[i] # [FEATURE_DIM] +gate[i] = softmax(gate_logit)[i] # [FEATURE_DIM] +out[b, k, i] = gate[i] * x[b, k, i] # [FEATURE_DIM] +``` + +**Tasks (bite-sized):** + +- Task 2D.1: VSN forward kernel (gate logits + softmax + scale) +- Task 2D.2: VSN backward kernel (softmax bwd + gate-weight grad + input grad) +- Task 2D.3: Numerical gradient check +- Task 2D.4: Wire VSN into PerceptionTrainer between snap-assemble and Mamba2 fwd +- Task 2D.5: Deploy fold-1 with VSN active — expect ≥ Phase 2C baseline + 0.5pt mean AUC + +**Acceptance criteria:** +- VSN gates softmax-normalise (numerical sum-check in test). +- Synthetic overfit converges (proves bwd chain is right). +- Fold-1 ≥ Phase 2C + 0.5pt mean AUC. +- **Diagnostic:** log VSN gates at epoch boundaries. Healthy gates show clear feature-importance differentiation across folds; degenerate VSN (all gates ≈ 1/FEATURE_DIM) indicates either too-strong regularisation or that the input features are too correlated for VSN to differentiate. + --- ## Phase 3: Attention Pool Over Mamba2 K-Positions