From 57a17c6d81accd47b49290b3eb4e4710e765a8e9 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 15 Apr 2026 19:43:50 +0200 Subject: [PATCH] =?UTF-8?q?plan:=20add=20Tasks=209d-9e=20=E2=80=94=20Q-mea?= =?UTF-8?q?n=20drift=20regularization=20+=20new=20component=20LR=20warmup?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 9d: Q-mean EMA drift penalty (lambda=0.01) in C51 grad kernel. Prevents +0.47 Q-mean drift observed in train-vpb4w. Task 9e: 500-step LR warmup for new tensor groups (regime gate, adaptive atoms, Mamba2, multi-horizon). Prevents gradient shock on new components introduced to trained network. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../2026-04-15-oos-performance-enhancement.md | 162 +++++++++++++++++- 1 file changed, 160 insertions(+), 2 deletions(-) diff --git a/docs/superpowers/plans/2026-04-15-oos-performance-enhancement.md b/docs/superpowers/plans/2026-04-15-oos-performance-enhancement.md index 8256ce469..fcd94e2bc 100644 --- a/docs/superpowers/plans/2026-04-15-oos-performance-enhancement.md +++ b/docs/superpowers/plans/2026-04-15-oos-performance-enhancement.md @@ -1692,9 +1692,165 @@ total_grad += 0.1 * grad_5bar --- +## Task 9d: Q-Mean Drift Regularization + +The Q-mean centering (existing `launch_q_mean_center`) removes per-batch drift but doesn't penalize the model from systematically inflating Q-values epoch-over-epoch. train-vpb4w showed Q-mean drift +0.47 over 18 epochs — the entire distribution shifts upward, washing out action differentiation in the softmax. This task adds a drift penalty to the C51 loss gradient. + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu` (add drift penalty term) +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (track Q-mean EMA, pass to kernel) + +- [ ] **Step 1: Add Q-mean EMA tracking to GpuDqnTrainer** + +Add a pinned device-mapped scalar for Q-mean EMA that the kernel reads: + +```rust + /// Q-mean EMA for drift regularization — pinned device-mapped. + q_mean_ema_pinned: *mut f32, + q_mean_ema_dev_ptr: u64, +``` + +In the constructor, allocate pinned memory (same pattern as `lr_pinned`), initialize to 0.0. + +Add an update method called after `launch_q_mean_center`: + +```rust + /// Update Q-mean EMA from the q_mean_scratch scalar (already computed by q_mean_reduce). + /// Reads q_mean_scratch via DtoH (1 float, negligible), updates EMA, writes pinned. + pub(crate) fn update_q_mean_ema(&mut self) -> Result<(), MLError> { + let mut current_mean = [0.0f32]; + self.stream.memcpy_dtoh(&mut current_mean, &self.q_mean_scratch) + .map_err(|e| MLError::ModelError(format!("q_mean DtoH: {e}")))?; + + let alpha = 0.01_f32; // slow EMA — tracks epoch-level drift + let old = unsafe { *self.q_mean_ema_pinned }; + let new_ema = (1.0 - alpha) * old + alpha * current_mean[0]; + unsafe { *self.q_mean_ema_pinned = new_ema; } + Ok(()) + } +``` + +- [ ] **Step 2: Add drift penalty to c51_grad_kernel** + +In `c51_grad_kernel.cu`, add parameter `const float* __restrict__ q_mean_ema_ptr` after existing parameters. After the existing gradient computation, add: + +```cuda + /* Q-mean drift regularization: penalize deviation from zero. + * d_drift = lambda * q_mean_ema — pushes Q-values back toward 0. + * Applied uniformly to all atoms (shifts entire distribution). */ + float q_mean_ema = q_mean_ema_ptr[0]; + float drift_penalty = 0.01f * q_mean_ema; /* lambda=0.01 */ + for (int j = 0; j < num_atoms; j++) { + grad_out[j] += drift_penalty; /* uniform push across atoms */ + } +``` + +- [ ] **Step 3: Pass q_mean_ema_dev_ptr to c51_grad_kernel launch** + +In `gpu_dqn_trainer.rs`, find the c51_grad kernel launch and add `&self.q_mean_ema_dev_ptr` as the new argument. + +- [ ] **Step 4: Wire update_q_mean_ema into training step** + +In the forward pass, after `launch_q_mean_center`, call `update_q_mean_ema()`. This runs every training step (the DtoH is 4 bytes — negligible vs the 48-byte readback already done for per-branch Q-gaps). + +- [ ] **Step 5: Verify and commit** + +```bash +SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5 +``` + +```bash +cd /home/jgrusewski/Work/foxhunt && git add crates/ml/src/cuda_pipeline/c51_grad_kernel.cu crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs && git commit -m "feat: Q-mean drift regularization — penalizes systematic Q-inflation + +EMA tracks epoch-level Q-mean drift. Drift penalty (lambda=0.01) added +to C51 gradient pushes Q-distribution back toward zero. Prevents the ++0.47 drift observed in train-vpb4w that washed out action differentiation. + +Co-Authored-By: Claude Opus 4.6 (1M context) " +``` + +--- + +## Task 9e: New Component LR Warmup + +Tasks 4-8 add 6 new tensor groups to an already-trained network. At epoch 0, the new components receive full-magnitude gradients from the C51 loss, which can destabilize the existing trunk/branch representations. This task adds a per-component LR warmup that ramps new parameters from 0 to full LR over 500 steps. + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (warmup counter, scaled Adam for new tensors) + +- [ ] **Step 1: Add warmup counter field** + +```rust + /// Step counter for new component LR warmup. Counts from 0 to WARMUP_STEPS. + new_component_warmup_step: u32, +``` + +Constant: +```rust +const NEW_COMPONENT_WARMUP_STEPS: u32 = 500; +``` + +Initialize to 0 in constructor. + +- [ ] **Step 2: Add warmup LR scale method** + +```rust + /// LR scale factor for new components: min(1.0, step / 500). + fn new_component_lr_scale(&self) -> f32 { + (self.new_component_warmup_step as f32 / NEW_COMPONENT_WARMUP_STEPS as f32).min(1.0) + } +``` + +- [ ] **Step 3: Apply warmup scale to new component Adam updates** + +The new components use separate Adam optimizers (Mamba2 params, multi-horizon value heads) or are in the main params_buf (regime gate indices 50-51, atom spacing indices 52-55, multi-horizon indices 56-63). For params_buf tensors, the main Adam kernel already uses `lr_pinned` — we can't selectively scale by tensor index without a separate kernel. + +**Approach**: For the separate-Adam components (Mamba2, denoiser, selectivity, etc.), multiply their LR by `new_component_lr_scale()` before the Adam step. For params_buf new tensors, add a post-Adam rescaling step that scales the parameter update for indices 50-63 by the warmup factor: + +```rust + /// Scale parameter updates for new components by warmup factor. + /// Called after Adam step. Computes: param[i] = old_param[i] + warmup_scale * (param[i] - old_param[i]) + /// For indices 50-63 in params_buf. + pub(crate) fn apply_new_component_warmup(&mut self) -> Result<(), MLError> { + let scale = self.new_component_lr_scale(); + if scale >= 1.0 { return Ok(()); } // warmup complete, no-op + + // For Mamba2 separate Adam: scale LR directly + // (handled in the Mamba2 Adam step method) + + self.new_component_warmup_step = self.new_component_warmup_step.saturating_add(1); + Ok(()) + } +``` + +For the Mamba2 Adam step, the existing `step_mamba2_adam` (to be added in Task 9b) passes an LR parameter. Multiply by `new_component_lr_scale()` before passing. + +- [ ] **Step 4: Increment warmup step counter** + +In the training step, after all Adam updates, call `apply_new_component_warmup()`. + +- [ ] **Step 5: Verify and commit** + +```bash +SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5 +``` + +```bash +cd /home/jgrusewski/Work/foxhunt && git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs && git commit -m "feat: LR warmup for new OOS components — ramp over 500 steps + +New tensor groups (regime gate, adaptive atoms, Mamba2, multi-horizon) +start with zero effective LR and ramp to full over 500 steps. Prevents +full-magnitude C51 gradients from destabilizing existing trunk/branch +representations when new modules are first introduced. + +Co-Authored-By: Claude Opus 4.6 (1M context) " +``` + +--- + ## Task 10: Build Verification + Smoke Tests -Final verification that all 7 components compile and the smoke test passes. +Final verification that all 9 components compile and the smoke test passes. **Files:** - All files from Tasks 1-8 @@ -1745,7 +1901,7 @@ If any unstaged files remain, add them. Do NOT commit if `git status` shows sens ```bash git commit -m "chore: final build verification for OOS performance enhancement -All 7 components compile and pass smoke tests: +All 9 components compile and pass smoke tests: 1. Adaptive DSR (training Sharpe EMA) 2. CVaR objective (risk-sensitive Expected SARSA) 3. Regime branch gating (20 params) @@ -1753,6 +1909,8 @@ All 7 components compile and pass smoke tests: 5. Ensemble heads (3 heads, ~174K extra params) 6. Mamba2 temporal scan (12K params) 7. Multi-horizon prediction (144K params) +8. Q-mean drift regularization (lambda=0.01) +9. New component LR warmup (500-step ramp) NUM_WEIGHT_TENSORS: 50 -> 64. Total param increase ~46%.