diff --git a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu index 201918af9..6d962fc2e 100644 --- a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu @@ -76,17 +76,17 @@ extern "C" __global__ void c51_grad_kernel( for (int dd = 0; dd < d; dd++) branch_offset += branch_sizes[dd] * num_atoms; - /* Magnitude gradient amplification: scale d==1 (magnitude) by 1.5× to - * compensate for Flat counterfactual discount (0.3× on ~58% of samples). - * Effective gradient per magnitude sample ≈ 0.3*0.58 + 1.0*0.42 = 0.594. - * Scale = 1.5 ≈ 1/0.594, conservative to avoid gradient explosion. - * Gradient clipping after this kernel prevents any runaway. */ - float mag_amp = (d == 1) ? 1.5f : 1.0f; + /* Per-branch loss weighting: magnitude (d==1) gets 0.2× C51 gradient. + * C51 cross-entropy inherently prefers low-variance actions (Small positions + * have tighter return distributions → lower cross-entropy). Reducing C51's + * influence on magnitude to 20% lets MSE (variance-neutral) dominate. + * Direction/order/urgency keep full C51 gradient for distributional risk awareness. */ + float branch_scale = (d == 1) ? 0.2f : 1.0f; - /* d_adv[b, d, a, j] = mag_amp * d_combined * (delta(a, a_d) - 1/A_d) */ + /* d_adv[b, d, a, j] = branch_scale * d_combined * (delta(a, a_d) - 1/A_d) */ for (int a = 0; a < A_d; a++) { float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A); - float grad_val = mag_amp * d_combined * dueling_grad; + float grad_val = branch_scale * d_combined * dueling_grad; int adv_idx = b * total_branch_atoms + branch_offset + a * num_atoms + j; atomicAdd(&d_adv_logits[adv_idx], grad_val); } diff --git a/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu b/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu index 8c0320cce..4e361692d 100644 --- a/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu @@ -79,17 +79,17 @@ extern "C" __global__ void mse_grad_kernel( for (int dd = 0; dd < d; dd++) branch_offset += branch_sizes[dd] * num_atoms; - /* Magnitude gradient amplification: scale d==1 (magnitude) by 1.5× to - * compensate for Flat counterfactual discount (0.3× on ~58% of samples). - * Effective gradient per magnitude sample ≈ 0.3*0.58 + 1.0*0.42 = 0.594. - * Scale = 1.5 ≈ 1/0.594, conservative to avoid gradient explosion. - * Gradient clipping after this kernel prevents any runaway. */ - float mag_amp = (d == 1) ? 1.5f : 1.0f; + /* Per-branch loss weighting: magnitude (d==1) gets 4.0× MSE gradient. + * MSE loss is variance-neutral (optimizes expected Q only, no distributional + * shape bias). Amplifying MSE for magnitude makes it the dominant loss signal, + * preventing C51's cross-entropy from collapsing magnitude to Small. + * Combined: magnitude = 0.2×C51 + 4.0×MSE ≈ 80% MSE-driven. */ + float branch_scale = (d == 1) ? 4.0f : 1.0f; - /* d_adv[b, d, a, j] = mag_amp * d_combined * (delta(a, a_d) - 1/A_d) */ + /* d_adv[b, d, a, j] = branch_scale * d_combined * (delta(a, a_d) - 1/A_d) */ for (int a = 0; a < A_d; a++) { float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A); - float grad_val = mag_amp * d_combined * dueling_grad; + float grad_val = branch_scale * d_combined * dueling_grad; int adv_idx = b * total_branch_atoms + branch_offset + a * num_atoms + j; atomicAdd(&d_adv_logits[adv_idx], grad_val); } diff --git a/docs/superpowers/plans/2026-04-08-per-branch-loss-weighting.md b/docs/superpowers/plans/2026-04-08-per-branch-loss-weighting.md new file mode 100644 index 000000000..a084f2e97 --- /dev/null +++ b/docs/superpowers/plans/2026-04-08-per-branch-loss-weighting.md @@ -0,0 +1,128 @@ +# Per-Branch C51/MSE Loss Weighting — Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Make the magnitude branch MSE-dominant (4× MSE, 0.2× C51) to prevent C51's variance-biased cross-entropy from collapsing magnitude to Small. + +**Architecture:** 2 CUDA kernel edits. Replace the existing `mag_amp = 1.5` with per-branch loss weighting: C51 grad scales magnitude by 0.2×, MSE grad scales magnitude by 4.0×. + +**Tech Stack:** CUDA (cubin via build.rs) + +--- + +### Task 1: Per-Branch Loss Weighting in Grad Kernels + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/mse_grad_kernel.cu` + +- [ ] **Step 1: Update C51 grad kernel** + +In `c51_grad_kernel.cu`, find the existing magnitude amplification block (around line 81): + +```c + /* Magnitude gradient amplification: scale d==1 (magnitude) by 1.5× to + * compensate for Flat counterfactual discount (0.3× on ~58% of samples). + * Effective gradient per magnitude sample ≈ 0.3*0.58 + 1.0*0.42 = 0.594. + * Scale = 1.5 ≈ 1/0.594, conservative to avoid gradient explosion. + * Gradient clipping after this kernel prevents any runaway. */ + float mag_amp = (d == 1) ? 1.5f : 1.0f; + + /* d_adv[b, d, a, j] = mag_amp * d_combined * (delta(a, a_d) - 1/A_d) */ + for (int a = 0; a < A_d; a++) { + float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A); + float grad_val = mag_amp * d_combined * dueling_grad; +``` + +Replace with: + +```c + /* Per-branch loss weighting: magnitude (d==1) gets 0.2× C51 gradient. + * C51 cross-entropy inherently prefers low-variance actions (Small positions + * have tighter return distributions → lower cross-entropy). Reducing C51's + * influence on magnitude to 20% lets MSE (variance-neutral) dominate. + * Direction/order/urgency keep full C51 gradient for distributional risk awareness. */ + float branch_scale = (d == 1) ? 0.2f : 1.0f; + + /* d_adv[b, d, a, j] = branch_scale * d_combined * (delta(a, a_d) - 1/A_d) */ + for (int a = 0; a < A_d; a++) { + float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A); + float grad_val = branch_scale * d_combined * dueling_grad; +``` + +- [ ] **Step 2: Update MSE grad kernel** + +In `mse_grad_kernel.cu`, find the same block (around line 80): + +```c + float mag_amp = (d == 1) ? 1.5f : 1.0f; + + /* d_adv[b, d, a, j] = mag_amp * d_combined * (delta(a, a_d) - 1/A_d) */ + for (int a = 0; a < A_d; a++) { + float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A); + float grad_val = mag_amp * d_combined * dueling_grad; +``` + +Replace with: + +```c + /* Per-branch loss weighting: magnitude (d==1) gets 4.0× MSE gradient. + * MSE loss is variance-neutral (optimizes expected Q only, no distributional + * shape bias). Amplifying MSE for magnitude makes it the dominant loss signal, + * preventing C51's cross-entropy from collapsing magnitude to Small. + * Combined: magnitude = 0.2×C51 + 4.0×MSE ≈ 80% MSE-driven. */ + float branch_scale = (d == 1) ? 4.0f : 1.0f; + + /* d_adv[b, d, a, j] = branch_scale * d_combined * (delta(a, a_d) - 1/A_d) */ + for (int a = 0; a < A_d; a++) { + float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A); + float grad_val = branch_scale * d_combined * dueling_grad; +``` + +- [ ] **Step 3: Verify build** + +```bash +SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5 +``` + +- [ ] **Step 4: Local smoke test** + +```bash +FOXHUNT_TEST_DATA=test_data/futures-baseline SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests::walk_forward --ignored --nocapture 2>&1 | grep -E 'Action div|Sharpe=|PF=' +``` + +Expected: diversity ≥5/9 at epoch 10. + +- [ ] **Step 5: Commit + push** + +```bash +git add crates/ml/src/cuda_pipeline/c51_grad_kernel.cu \ + crates/ml/src/cuda_pipeline/mse_grad_kernel.cu +git commit -m "feat(magnitude): per-branch loss weighting — 0.2× C51 + 4.0× MSE for magnitude (variance-neutral)" +git push origin main +``` + +--- + +### Task 2: H100 Validation + +- [ ] **Step 1: Delete old fxcache + deploy** + +```bash +# Delete cached fxcache on PVC (force rebuild with new binary) +kubectl run cleanup --image=ubuntu:24.04 --restart=Never -n foxhunt \ + --overrides='...' # (delete /feature-cache/*.fxcache) + +./scripts/argo-train.sh dqn --baseline --epochs 10 +``` + +- [ ] **Step 2: Monitor** + +Watch for: +- Diversity ≥5/9 at epochs 5-10 (not 3/9) +- Sharpe > 0 sustained +- PF ≥ 1.0 + +```bash +argo logs -n foxhunt --follow=false -c main 2>&1 | grep 'Action div' +``` diff --git a/docs/superpowers/specs/2026-04-08-per-branch-loss-weighting-design.md b/docs/superpowers/specs/2026-04-08-per-branch-loss-weighting-design.md new file mode 100644 index 000000000..0442223c9 --- /dev/null +++ b/docs/superpowers/specs/2026-04-08-per-branch-loss-weighting-design.md @@ -0,0 +1,49 @@ +# Per-Branch C51/MSE Loss Weighting — Design Spec + +## Problem + +C51 distributional cross-entropy loss collapses the magnitude branch to Small (lowest variance) within 4 epochs. MSE loss doesn't have this bias — magnitude diversity holds at 7/9 during MSE warmup epochs. When C51 takes over, diversity drops to 3/9. + +Root cause: C51 optimizes distributional shape (cross-entropy). Small positions produce tighter return distributions → lower cross-entropy → C51 structurally prefers Small. MSE optimizes expected value only (scalar) → no variance bias. + +## Solution + +Per-branch gradient scaling in the C51 and MSE grad kernels. The magnitude branch (d==1) gets reduced C51 gradient (0.2×) and amplified MSE gradient (4.0×), making it MSE-dominant. All other branches (direction, order, urgency) are unchanged. + +### Effective loss blend per branch + +| Branch | C51 weight | MSE weight | Dominant | +|--------|-----------|-----------|----------| +| Direction (d=0) | 1.0× (normal) | 1.0× (normal) | C51 | +| **Magnitude (d=1)** | **0.2×** | **4.0×** | **MSE** | +| Order (d=2) | 1.0× (normal) | 1.0× (normal) | C51 | +| Urgency (d=3) | 1.0× (normal) | 1.0× (normal) | C51 | + +### Why this works + +- MSE gradient `d_combined = isw * td_error * p_j * (z_j - E[Q])` treats Q as a scalar expectation. No distributional shape bias. Magnitude can differentiate based on expected return alone. +- C51 gradient `d_combined = isw * (exp(lp) - proj)` optimizes the full distribution. Keeps distributional risk awareness for direction/order/urgency. +- The 0.2×/4.0× factors replace the existing `mag_amp = 1.5` from the earlier gradient amplification fix. The 1.5× is subsumed — magnitude now gets 4.0× MSE (which is stronger than 1.5× undifferentiated). + +## Files Modified + +- `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu` — scale `d==1` gradient by 0.2× +- `crates/ml/src/cuda_pipeline/mse_grad_kernel.cu` — scale `d==1` gradient by 4.0× + +## Files NOT Modified + +- Loss kernels (c51_loss, mse_loss) — unchanged. The per-sample loss for PER priorities is unaffected. +- Rust launch code — no new parameters. +- CUDA graph capture — no buffer size changes. +- IQN head — operates independently on its own gradient path. + +## Testing + +- Smoke test: magnitude diversity should hold ≥5/9 through epoch 10 (not collapse at epoch 4) +- Q-value gap should remain >0 (direction differentiation preserved) +- PF should remain ≥1.0 (profitability preserved) + +## Success Criteria + +- Magnitude diversity ≥5/9 at epoch 10 on H100 (was 3/9 with uniform weighting) +- No regression in Sharpe or PF vs current baseline (Sharpe=1.15, PF=1.22 peak)