feat(magnitude): per-branch loss weighting — 0.2× C51 + 4.0× MSE for magnitude
C51 cross-entropy structurally prefers low-variance actions (Small positions have tighter return distributions). MSE is variance-neutral. By reducing C51's gradient to 20% and amplifying MSE to 4× for the magnitude branch, MSE becomes the dominant loss signal for position sizing. Direction/order/urgency keep full C51 gradient for distributional risk awareness. Replaces the previous mag_amp=1.5 which addressed a symptom (Flat gradient starvation) rather than the root cause (C51 variance bias). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
128
docs/superpowers/plans/2026-04-08-per-branch-loss-weighting.md
Normal file
128
docs/superpowers/plans/2026-04-08-per-branch-loss-weighting.md
Normal file
@@ -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 <workflow> -n foxhunt --follow=false -c main 2>&1 | grep 'Action div'
|
||||
```
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user