# 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)
