plan(dqn-v2): Plan 4 full-scope reality reconciliation
Comprehensive revision to match landed Plan 1+2+3 codebase state. Pattern
mirrors the Plan 3 second revision: per-task "Reality reconciliation"
blocks at the top of each task body identifying what's stale vs. landed,
with concrete file paths in the actual cuda_pipeline tree.
Added:
- Task dependency graph with recommended execution order:
5 (light ISV) → 4 (refactor) → 2 (GRN ADOPT) → 1 (Full VSN) →
6 (aux heads) → 3 (multi-Q IQN) → 7 (audit) → 8 (Argo)
- Per-task implementation surface with concrete trunk hooks:
- Task 1: pre-h_s1 VSN gate in batched_forward.rs::forward_online_raw
- Task 2: GRN audit confirms ADOPT branch (GRN absent from DQN trunk;
only in ml-supervised TFT). Replaces h_s1/h_s2 Linear blocks.
- Task 3: re-scoped — IQN already runs num_quantiles=32 with random τ.
Task is CONSTRAINING to fixed τ ∈ {.05, .25, .5, .75, .95}.
- Task 4: pure Rust API split (no kernel changes, no checkpoint break)
around existing forward_online_raw structure
- Task 5: split into Mode A (light, pre-Task-1, 3 ISV slots) and
Mode B (full, post-Task-1, 7 ISV slots). Mode A recommended first.
- Task 6: aux head loss scaled by ISV[LEARNING_HEALTH] per pearl
- Checkpoint-break consolidation note: Tasks 2/1/6/3 each break checkpoint
compatibility; land in sequence with no Argo run between (one fingerprint
shift per commit; final Argo at Task 8 amortizes retraining cost)
- Task 8 absorbs Plan 3's deferred Argo Tier 1 gate (combined validation)
No code changes.
This commit is contained in:
jgrusewski
@@ -93,18 +93,75 @@ Read `docs/ml-supervised-to-dqn-concept-audit.md` (delivered at end of Plan 1).
|
||||
|
||||
If the audit doc does NOT yet have a GRN decision: stop, return to Plan 1 Task 6 (orphan audit) or run the A.5 audit to produce one. Do not proceed with a missing decision.
|
||||
|
||||
**GRN audit decision (2026-04-25 verification):** The DQN trunk in `crates/ml/src/cuda_pipeline/batched_forward.rs::forward_online_raw` is `Linear→ReLU→Linear` (no LayerNorm, no ELU, no GLU gate). GRN exists in `crates/ml-supervised::tft::gated_residual` but is NOT wired into the DQN trunk. **Decision: Branch A (ADOPT)** for Task 2.
|
||||
|
||||
---
|
||||
|
||||
## Task dependency graph + recommended execution order
|
||||
|
||||
```dot
|
||||
digraph plan4_deps {
|
||||
rankdir=LR;
|
||||
"5 (attn ISV)" [shape=box style=filled fillcolor=lightgreen];
|
||||
"7 (audit)" [shape=box style=filled fillcolor=lightgreen];
|
||||
"2 (GRN)" [shape=box];
|
||||
"4 (E/D split)" [shape=box];
|
||||
"1 (Full VSN)" [shape=box style=filled fillcolor=lightyellow];
|
||||
"3 (multi-Q)" [shape=box style=filled fillcolor=lightyellow];
|
||||
"6 (aux heads)" [shape=box];
|
||||
"8 (validation)"[shape=box style=filled fillcolor=lightblue];
|
||||
|
||||
"5 (attn ISV)" -> "1 (Full VSN)" [label="diagnostic"];
|
||||
"5 (attn ISV)" -> "6 (aux heads)" [label="diagnostic"];
|
||||
"2 (GRN)" -> "4 (E/D split)" [label="trunk shape"];
|
||||
"4 (E/D split)"-> "1 (Full VSN)" [label="encoder boundary"];
|
||||
"4 (E/D split)"-> "3 (multi-Q)" [label="decoder boundary"];
|
||||
"4 (E/D split)"-> "6 (aux heads)"[label="decoder peer"];
|
||||
"1 (Full VSN)" -> "8 (validation)";
|
||||
"2 (GRN)" -> "8 (validation)";
|
||||
"3 (multi-Q)" -> "8 (validation)";
|
||||
"6 (aux heads)"-> "8 (validation)";
|
||||
"7 (audit)" -> "8 (validation)";
|
||||
}
|
||||
```
|
||||
|
||||
**Recommended order:**
|
||||
1. **Task 5** — ISV-only attention diagnostics (light, no checkpoint break, gives telemetry for the heavier tasks)
|
||||
2. **Task 4** — Encoder/decoder Rust-API split (refactor only, no kernel changes, no checkpoint break — sets clean boundaries for 1/3/6)
|
||||
3. **Task 2** — GRN ADOPT (real architectural change, but localized to the 2 trunk Linear blocks; checkpoint break)
|
||||
4. **Task 1** — Full VSN over feature groups (depends on E/D boundary; checkpoint break)
|
||||
5. **Task 6** — Aux heads (depends on E/D boundary; checkpoint break)
|
||||
6. **Task 3** — Multi-quantile IQN (most invasive — touches all IQN consumers; checkpoint break)
|
||||
7. **Task 7** — Audit close-out (pure docs, after 1-6)
|
||||
8. **Task 8** — Argo validation run (after 1-7 land)
|
||||
|
||||
**Checkpoint break consolidation:** Tasks 2, 1, 6, 3 each break checkpoint compatibility. Land them in sequence with NO Argo run between (one fingerprint shift per commit; final Argo run at Task 8 amortizes the retraining cost across all changes). Revert to a known-good commit if any single task regresses the local smoke beyond ~25-point Sharpe drop.
|
||||
|
||||
---
|
||||
|
||||
## Task 1: E.1 Full TFT Variable Selection Network
|
||||
|
||||
**Current state:** Two partial VSN masks exist — `vsn_mag`, `vsn_dir` — covering the direction and magnitude branches only. Apply selectively. The spec requires VSN over all state feature groups: market features, OFI features, MTF features, portfolio state, plan_isv block, TLOB features (post-Plan 2 D.8), and Mamba2 temporal state.
|
||||
**Reality reconciliation (2026-04-25).** The current `vsn_mag` / `vsn_dir` are **scalars** (per-batch mean of |w| across each branch's value-head linear) computed in `training_loop.rs:2425`, not per-feature-group masks. They expose existing weight magnitudes for diagnostic display, NOT a learned per-group selector. Task 1 is a real model expansion: add 6 learnable softmax-over-groups gating layers operating on the assembled state vector AT the encoder input (before `forward_online_raw`'s `h_s1` GEMM). Each group gets a tiny per-group GLU-style projection: `mask_g = softmax_over_groups( Linear_g(group_features_g) )`. The state vector then becomes `state_gated = state * mask_broadcast_per_group`.
|
||||
|
||||
**Files:**
|
||||
- Create: `crates/ml/src/cuda_pipeline/vsn_feature_selection_kernel.cu`
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — allocate VSN weights per feature group, pre-trunk hook
|
||||
**Implementation surface (concrete):**
|
||||
- **Encoder hook**: `crates/ml/src/cuda_pipeline/batched_forward.rs::forward_online_raw` line 455+. Insert a pre-h_s1 VSN gate that mutates the input state buffer in place (or writes to a separate gated-state scratch, then h_s1 uses that).
|
||||
- **VSN params**: 6 small `[group_dim, group_dim]` projections + 6 bias vectors. Total params ≈ Σ group_dim² ≈ 42² + 32² + 16² + 16² + 8² + 7² ≈ 3,250 floats. Tiny relative to existing trunk (~250k params).
|
||||
- **Adam state**: 2× param count for Adam moments. Add to `gpu_dqn_trainer.rs` parameter buffer block.
|
||||
- **Forward kernel** `vsn_feature_selection_kernel.cu`: per-group projection + softmax-over-groups + gate-into-input. Use cuBLAS for the GEMMs; one custom kernel for the softmax-over-groups + multiply.
|
||||
- **Backward kernel**: gradient through softmax + per-group projection. Standard pattern.
|
||||
- **Feature-group ranges**: already defined in `crates/ml-core/src/state_layout.rs` (MARKET_START=0, OFI_START=42, TLOB_START=74, MTF_START=90, PORTFOLIO_START=106, PLAN_ISV_START=114, padding from 121 — note STATE_DIM_PADDED=128).
|
||||
- **Diagnostic**: emit per-group mean-mask as 6 GPU-written ISV slots (Task 5 may also expose these).
|
||||
|
||||
**Risk:** First commit that adds non-trivial model parameters. If grad propagation breaks, the pearl-style debugging is well-trodden in this codebase (Plan 2 Task 2 grad-check pattern). Run grad-check smoke before declaring done.
|
||||
|
||||
**Files (concrete):**
|
||||
- Create: `crates/ml/src/cuda_pipeline/vsn_feature_selection_kernel.cu` (forward + backward)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs::forward_online_raw` (pre-h_s1 hook)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/batched_backward.rs` (mirror gradient path)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (param + Adam state allocation, ISV slot allocation for 6 group masks)
|
||||
- Modify: `crates/ml/build.rs` — register new kernel
|
||||
- Modify: `crates/ml/src/cuda_pipeline/state_layout.cuh` — expose feature-group index ranges
|
||||
- Test: `vsn_feature_selection_produces_nonzero_masks`
|
||||
- Modify: `crates/ml/src/cuda_pipeline/state_layout.cuh` — alias the existing OFI/TLOB/MTF/PORTFOLIO/PLAN_ISV ranges as `SL_*_GROUP_BEGIN/END` macros (they already exist as `SL_*_START`)
|
||||
- Test: rely on the standard `multi_fold_convergence` smoke; add a grad-check test for VSN params using the Plan 2 Task 2 grad-check helper if one exists, otherwise skip the bespoke test.
|
||||
|
||||
### Subtask 1A: Define feature-group index ranges
|
||||
|
||||
@@ -326,9 +383,28 @@ git commit -m "plan4(task1C): upgrade E.1 VSN to per-feature weights (cond. on 1
|
||||
|
||||
---
|
||||
|
||||
## Task 2: E.2 Gated Residual Network (conditional: ADOPT vs CANONICALISE)
|
||||
## Task 2: E.2 Gated Residual Network — Branch A (ADOPT) confirmed
|
||||
|
||||
**Decision gate:** `docs/ml-supervised-to-dqn-concept-audit.md` row for GRN.
|
||||
**Reality reconciliation (2026-04-25).** GRN audit completed: GRN exists in `crates/ml-supervised/src/tft/gated_residual.rs` (Rust + supervised pretraining only) and `crates/ml/src/transformers/temporal_fusion.rs` (TFT inference path). It is **NOT** wired into the DQN cuda_pipeline trunk. The DQN trunk in `batched_forward.rs::forward_online_raw` is `Linear→ReLU→Linear` with no LayerNorm, no ELU, no GLU. **Branch A (ADOPT) applies.** Branch B is moot.
|
||||
|
||||
**Implementation surface (concrete):**
|
||||
- **Trunk site**: `batched_forward.rs::forward_online_raw` lines 471–481 (h_s1 + h_s2 GEMMs). Replace each `Linear→ReLU→Linear` block with a GRN block.
|
||||
- **GRN block formula** (per spec §4.E.2): `GRN(x) = LayerNorm(x + GLU(Linear_b(ELU(Linear_a(x)))))` with `Linear_a: SH1×SD`, `Linear_b: 2*SH1 × SH1` (the 2× is for the GLU split), `LayerNorm` on `[SH1]`. Skip the optional context input `c` for trunk use.
|
||||
- **New params per GRN**: 1 LayerNorm γ/β + 2 Linear layers + bias. For h_s1 GRN: ~2 × SH1 × SD + 2*SH1 + 2 × 2*SH1 × SH1 + 2*SH1 + 2*SH1 (LN) ≈ 80k floats with SH1=128, SD=128. For h_s2 GRN similar. Total Adam state 4×.
|
||||
- **Forward kernel** `grn_kernel.cu`: (a) Linear_a via cuBLAS, (b) ELU element-wise, (c) Linear_b via cuBLAS (output 2*SH1), (d) GLU split (chunk[0] * sigmoid(chunk[1])) custom kernel, (e) residual add + LayerNorm fused.
|
||||
- **Backward kernel**: gradient through LayerNorm → residual → GLU → Linear_b → ELU → Linear_a. Six sub-steps, each well-known. Use existing `cuda_autograd` pattern.
|
||||
- **Checkpoint break**: yes — replaces Linear weights with GRN sub-block weights of different shape. Layout fingerprint already shifts on any param-block size change.
|
||||
|
||||
**Files (concrete):**
|
||||
- Create: `crates/ml/src/cuda_pipeline/grn_kernel.cu` (forward + backward + LayerNorm + GLU helpers)
|
||||
- Create: `crates/ml/src/cuda_pipeline/gpu_grn.rs` — Rust wrapper holding GRN params + Adam state
|
||||
- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs` — swap h_s1/h_s2 calls
|
||||
- Modify: `crates/ml/src/cuda_pipeline/batched_backward.rs` — mirror gradient path
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — param + Adam allocation, replace W_s1/W_s2 entries
|
||||
- Modify: `crates/ml/build.rs` — register new kernel
|
||||
- Modify: `docs/ml-supervised-to-dqn-concept-audit.md` — flip GRN row to "WIRED" with commit SHA
|
||||
|
||||
**Decision gate (historical):** `docs/ml-supervised-to-dqn-concept-audit.md` row for GRN.
|
||||
|
||||
### Branch A: ADOPT (GRN not present in DQN trunk)
|
||||
|
||||
@@ -393,12 +469,22 @@ git commit -m "plan4(task2B): E.2 CANONICALISE — existing GRN impl verified, a
|
||||
|
||||
## Task 3: E.3 Multi-quantile IQN decomposition
|
||||
|
||||
**Current state:** IQN branch produces one CVaR-weighted Q estimate. Spec wants explicit multi-quantile heads at 5/25/50/75/95 quantiles.
|
||||
**Reality reconciliation (2026-04-25).** The original task description was misleading. The DQN already has a multi-quantile IQN: `gpu_iqn_head.rs` line 90 sets `num_quantiles: 32` (standard IQN paper, random τ per batch). The plan's "extend IQN to 5/25/50/75/95 quantile decomposition" is actually a **CONSTRAINT** swap — replace random τ ∈ U[0,1] sampling with FIXED τ ∈ {0.05, 0.25, 0.50, 0.75, 0.95} so the 5 quantile heads are interpretable / consistent across batches. This is much lighter than I initially expected: the kernel structure stays the same; only the τ sampler changes.
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/iqn_kernel.cu` — emit 5 quantile estimates instead of one
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — extend IQN output buffers
|
||||
- Test: `iqn_multi_quantile_heads_produce_monotonic_estimates`
|
||||
**Implementation surface (concrete):**
|
||||
- **Config**: change `num_quantiles: 32 → 5` in `gpu_iqn_head.rs::Config::default()`. Add `pub fixed_taus: bool` flag (default true post-Task 3) so we can toggle back to random for comparison.
|
||||
- **τ sampler**: in `iqn_dual_head_kernel.cu`, replace the per-call random τ with a host-passed `fixed_taus[5]` array when the flag is set. cuBLAS dimension shrinks from `[B × 32]` to `[B × 5]` on the quantile axis — substantial buffer-size reduction.
|
||||
- **Loss kernel**: existing IQN quantile-Huber regression handles arbitrary num_quantiles. No formula change.
|
||||
- **Q-greedy aggregation**: where IQN currently averages all 32 quantiles for action selection, switch to using the median (τ=0.50) for action selection and store the other 4 quantiles for CVaR/risk diagnostics.
|
||||
- **Diagnostics**: 4 new ISV slots tail-appended to expose Q_p05/Q_p25/Q_p75/Q_p95 EMAs (Q_p50 is the existing greedy Q). These are GPU-written each epoch.
|
||||
- **Checkpoint break**: yes — IQN head output dimension changes (32→5).
|
||||
|
||||
**Files (concrete):**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu` — fixed-τ sampler branch
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` — config, buffer sizing, action-selection median path
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — 4 new ISV slots, smaller IQN buffer alloc
|
||||
- Modify: `crates/ml/src/cuda_pipeline/iqn_cvar_kernel.cu` — CVaR-weighted reduction now uses fixed quantiles
|
||||
- Test: `multi_fold_convergence` smoke; verify Sharpe doesn't regress > 25 points (5 quantiles is much less than 32 — some signal loss expected; budget for ≤15 point Sharpe reduction)
|
||||
|
||||
- [ ] **Step 3.1: Write failing test**
|
||||
|
||||
@@ -448,13 +534,21 @@ git commit -m "plan4(task3): E.3 multi-quantile IQN heads (5/25/50/75/95)"
|
||||
|
||||
## Task 4: E.4 Encoder-Decoder separation
|
||||
|
||||
**Current state:** single "trunk" function. Spec wants an explicit `StateEncoder` + per-branch `ValueDecoder`.
|
||||
**Reality reconciliation (2026-04-25).** Original plan said "rename `trunk_forward.rs` → `state_encoder.rs`". That file does not exist — the trunk forward lives in `crates/ml/src/cuda_pipeline/batched_forward.rs::forward_online_raw`. The function already has the encoder/decoder shape internally (h_s1, h_s2 are encoder activations; h_v + 4 branch heads are decoders) but does not expose this as separate Rust types. Task 4 is **pure Rust API refactor** — extract two struct boundaries from the existing single function. **No kernel changes, no checkpoint break, no parameter movement.** It exists to give Tasks 1, 3, 6 clean attachment points.
|
||||
|
||||
**Files:**
|
||||
- Rename: `crates/ml/src/cuda_pipeline/trunk_forward.rs` → `state_encoder.rs`
|
||||
- Create: `crates/ml/src/cuda_pipeline/value_decoder.rs`
|
||||
- Modify: `gpu_dqn_trainer.rs` — route through encoder + 4 decoder branches
|
||||
- Test: `encoder_decoder_boundary_is_explicit`
|
||||
**Implementation surface (concrete):**
|
||||
- **Extract `StateEncoder`**: a thin wrapper around the existing trunk GEMMs `h_s1` + `h_s2`. Owns the trunk Linear weights (post-Task 2 these become GRN sub-blocks). Exposes `forward(state) -> encoded_h_s2_handle`.
|
||||
- **Extract `ValueDecoder`**: 4 branch decoders (Dir, Mag, Ord, Urg), each owning the per-branch heads + IQN/C51 ensemble. Exposes `forward(encoded) -> BranchDecoderOutput`.
|
||||
- **`BranchDecoderOutput`**: existing `q_per_action`, plus the V_short/V_long decomposition Plan 2 D.3 already landed.
|
||||
- **Plan 2 D.3 integration**: V_short/V_long already in IQL value head; just expose them on the decoder struct.
|
||||
|
||||
**Risk**: cross-cutting refactor touches every call site that invokes the trunk forward. Mechanical — no behavior change. Compile errors will catch most issues.
|
||||
|
||||
**Files (concrete):**
|
||||
- Create: `crates/ml/src/cuda_pipeline/state_encoder.rs` (Rust wrapper around existing trunk GEMMs)
|
||||
- Create: `crates/ml/src/cuda_pipeline/value_decoder.rs` (Rust wrapper around existing per-branch heads)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs::forward_online_raw` — split into encoder/decoder method calls (or keep as one function and add `encoder_part` + `decoder_part` helpers)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — wire encoder/decoder calls through the existing infrastructure
|
||||
|
||||
- [ ] **Step 4.1: Write failing test**
|
||||
|
||||
@@ -528,6 +622,31 @@ git commit -m "plan4(task4): E.4 encoder-decoder separation, explicit Branch dec
|
||||
|
||||
## Task 5: E.5 Attention-weight interpretability via ISV
|
||||
|
||||
**Reality reconciliation (2026-04-25).** Original plan allocated slots in [63..72) which now collide with Plan 3 Task 1 reward-component EMAs at [63..69). **Tail-append at current ISV_TOTAL_DIM = 87.** The slot count also depends on Task 1 (Full VSN) — if E.1 lands first, Task 5 exposes 6 per-group VSN masks; if E.5 runs before E.1, only the existing 2 scalar `vsn_mag`/`vsn_dir` plus a Mamba2 retention proxy can be exposed (3 slots).
|
||||
|
||||
**Two execution modes:**
|
||||
- **Mode A (recommended, light, no-blocker):** Task 5 runs BEFORE Task 1. Exposes 3 ISV slots: `VSN_MAG_EMA`, `VSN_DIR_EMA`, `MAMBA2_RETENTION_EMA`. Pure ISV addition. ~150 LOC. Diagnostic only.
|
||||
- **Mode B (full, post-Task-1):** Task 5 runs AFTER Task 1. Exposes 7 slots: 6 per-group VSN masks + Mamba2 retention. Producer reads VSN forward kernel output buffers.
|
||||
|
||||
Choose Mode A as the default. Mode B is upgrade path if E.1 lands and exposes per-group masks.
|
||||
|
||||
**Implementation surface — Mode A (concrete):**
|
||||
- **`VSN_MAG_EMA` / `VSN_DIR_EMA`**: existing scalars in `training_loop.rs:2425` are CPU-computed per epoch. Migrate to GPU-side EMA: add a tiny `vsn_focus_ema_kernel.cu` that takes the per-batch mean |w| of branch value-head weights (already computed) and EMAs into the slot. Adaptive α matching Plan 3 Task 3 convention.
|
||||
- **`MAMBA2_RETENTION_EMA`**: Mamba2 has no attention; use `mean(|state_transition|)` as a retention proxy. Mamba2 forward kernel emits a per-batch scalar to a pinned buffer; tiny reduction EMAs into ISV slot.
|
||||
- **Slot tail-append**: 3 new slots at ISV[87..90). Fingerprint shifts [85,86] → [90, 91]. ISV_TOTAL_DIM 87 → 92.
|
||||
- **Read-only monitor**: `AttentionMonitor` mirrors `PlanThresholdMonitor` pattern.
|
||||
|
||||
**Files (Mode A, concrete):**
|
||||
- Create: `crates/ml/src/cuda_pipeline/attention_focus_ema_kernel.cu` (single tiny kernel)
|
||||
- Create: `crates/ml/src/trainers/dqn/monitors/attention_monitor.rs`
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — 3 ISV slots + fingerprint shift
|
||||
- Modify: `crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu` — emit retention scalar to pinned scratch (single line)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` — launch the EMA kernel alongside Plan 3 Task 1's reward-EMA launcher
|
||||
|
||||
### Original task body (Mode B, blocked on E.1)
|
||||
|
||||
The body below assumes E.1 has landed. Use as the post-E.1 expansion plan if/when Task 1 lands.
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — 9 new ISV slots [63..72)
|
||||
|
||||
@@ -631,13 +750,24 @@ git commit -m "plan4(task5): E.5 attention-weight ISV interpretability (slots [6
|
||||
|
||||
## Task 6: E.6 Multi-task auxiliary heads
|
||||
|
||||
**Current state:** single-objective DQN (policy Q). Spec adds two auxiliary heads: next-bar return prediction + 5-bar regime classification. Aux losses are small and feed gradients to the trunk only (heads are lightweight).
|
||||
**Reality reconciliation (2026-04-25).** Plan calls for two aux heads: (a) next-bar return regression, (b) 5-bar regime classification. The targets exist already: regime_stability is at `isv_signals[11]` (current bar's regime stability EMA), and next_bar return is the existing forward-return label used by the supervised path. The auxiliary heads are small MLPs branching off the encoder output (h_s2). Per pearl_one_unbounded_signal_per_reward, aux loss should be SCALED by a small bounded multiplier (typical aux-loss weight ~0.1-0.3) applied to the standard MSE/CE loss; combined with shaping_scale and learning_health, no tuned coefficient blows up.
|
||||
|
||||
**Files:**
|
||||
- Create: `crates/ml/src/cuda_pipeline/aux_heads_kernel.cu`
|
||||
- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs` — compute aux loss, add to total loss
|
||||
- Modify: `gpu_dqn_trainer.rs` — allocate aux head weights
|
||||
- Test: `aux_heads_produce_meaningful_signal`
|
||||
**Implementation surface (concrete):**
|
||||
- **Aux Head 1 (next-return regression)**: `Linear(SH2 → 1)`. Loss: MSE against `next_close_pct = (raw_next - raw_close) / raw_close`. Already in `experience_env_step`'s scope.
|
||||
- **Aux Head 2 (regime classification)**: `Linear(SH2 → K_regime_classes)`. Spec says 5-bar regime classification. Need to discretize `isv_signals[11]` into K bins (e.g., low/mid/high stability = 3 classes is simpler than 5). Loss: cross-entropy against the bucketed label.
|
||||
- **Aux loss combination**: `aux_loss = aux_w * (MSE_return + CE_regime)` where `aux_w = aux_base × ISV[LEARNING_HEALTH_INDEX]` (gates by health — aux signal trusted more as training matures). `aux_base ≈ 0.1` is the only tuned constant (small fraction of the policy loss).
|
||||
- **Backward path**: aux gradients flow into encoder via existing trunk backward. Aux heads themselves get their own gradient. No interaction with C51/IQN.
|
||||
- **Diagnostic**: ISV slots for aux loss EMAs (1 for return MSE, 1 for regime CE accuracy) — tail-appended after Task 5's slots.
|
||||
- **Checkpoint break**: yes (new heads + Adam state).
|
||||
|
||||
**Files (concrete):**
|
||||
- Create: `crates/ml/src/cuda_pipeline/aux_heads_kernel.cu` — two small forward + backward kernels
|
||||
- Create: `crates/ml/src/cuda_pipeline/gpu_aux_heads.rs` — Rust wrapper holding aux params + Adam state
|
||||
- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs` — emit `h_s2` to aux-heads forward (in parallel with branch decoders)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/batched_backward.rs` — accumulate aux gradient into `h_s2` backward
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — param + Adam, ISV slot allocation
|
||||
- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs` — compute aux loss EMA, log via HEALTH_DIAG
|
||||
- Modify: `crates/ml/build.rs` — register new kernel
|
||||
|
||||
- [ ] **Step 6.1: Write failing test**
|
||||
|
||||
@@ -733,6 +863,8 @@ git commit -m "plan4(task6): E.6 aux heads (next-bar return MSE + 5-bar regime C
|
||||
|
||||
## Task 7: Part E audit close-out
|
||||
|
||||
**Reality reconciliation (2026-04-25).** Pure docs task. Set each Part E audit row to one of `LANDED <SHA>`, `LANDED-via-Plan-2 <SHA>`, or `OUT-confirmed`. No code changes. Run after Tasks 1–6 land; before Task 8 (validation).
|
||||
|
||||
**Files:**
|
||||
- Modify: `docs/ml-supervised-to-dqn-concept-audit.md` — update all E-rows to their landed state
|
||||
|
||||
@@ -771,13 +903,16 @@ git commit -m "plan4(task7): Part E audit close-out — every supervised concept
|
||||
|
||||
## Task 8: Plan 4 validation run — multi-seed × multi-fold, Tier 1 retention check
|
||||
|
||||
**Reality reconciliation (2026-04-25).** Argo L40S deploy. Combines Plan 3's deferred Argo run + Plan 4's retention check into one combined validation run since Plan 3 Task 10 was deferred. Multi-seed × 60-epoch run with the full Plan 1-4 substrate landed.
|
||||
|
||||
Plan 4 is a concept-adoption plan. Its value is measured over time (richer representations → better Tier 3 profitability later). Plan 4 exit only requires:
|
||||
|
||||
1. Tier 1 convergence is RETAINED (no regression from Plan 3).
|
||||
2. Aux heads produce meaningful signal (Task 6 smoke passes in live run).
|
||||
3. All Task 1–7 smoke tests pass.
|
||||
1. Tier 1 convergence is RETAINED (no regression from Plan 3 baseline that hasn't been Argo-validated yet — this is the first opportunity to confirm).
|
||||
2. Aux heads produce meaningful signal (Task 6 monitor shows non-trivial regime-classification accuracy + return-MSE decline).
|
||||
3. All Task 1–7 local smoke tests pass.
|
||||
4. ISV signal monitors show expected behavior: VSN per-group masks (Task 1) > 0; per-quantile Q values (Task 3) approximately monotonic.
|
||||
|
||||
Tier 2 and 3 exit gates are checked by Plan 5.
|
||||
Tier 2 and 3 exit gates are checked by Plan 5. Tier 1 gate (and Plan 3 retroactive Tier 1 + Tier 2 behavioural gates) ALL pass here.
|
||||
|
||||
- [ ] **Step 8.1: Update metric-bands.toml**
|
||||
|
||||
|
||||
Reference in New Issue
Block a user