spec: direction-conditioned magnitude head (Layer 3)

Magnitude FC input changes from h_s2[B, SH2] to [h_s2; Q_dir][B, SH2+3]. Direction E[Q] values (3 scalars per sample) are computed from branch 0 logits and concatenated. w_b1fc grows by 3*adv_h parameters (768 for AH=256). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-14 22:14:49 +02:00
parent b2eb1b7741
commit 234870d966
1 changed files with 146 additions and 0 deletions
--- a/docs/superpowers/specs/2026-04-14-direction-conditioned-magnitude-design.md
+++ b/docs/superpowers/specs/2026-04-14-direction-conditioned-magnitude-design.md
@@ -0,0 +1,146 @@
+# Direction-Conditioned Magnitude Head — Design Spec
+
+**Goal:** Let the magnitude branch see what direction was chosen, so Long+Full and Flat+Full have different Q-values.
+
+**Problem:** The 4-branch DQN treats branches as independent. Each branch FC reads the same `h_s2[B, SH2]` trunk. The magnitude head cannot distinguish Long+Full (aggressive, high risk) from Flat+Full (nonsensical, no direction). With 3 direction actions, the magnitude head sees 3 identical inputs and must produce 3 × 51 logits without knowing which direction was selected.
+
+**Architecture:** Feed direction E[Q] as extra input to the magnitude FC layer. Only the magnitude branch (branch 1) changes; direction, order, and urgency branches remain unchanged.
+
+---
+
+## Current Forward Pass
+
+```
+h_s2[B, SH2]  ──┬── W_b0fc[AH, SH2] ──→ h_b0[B, AH] ──→ W_b0out ──→ dir_logits[B, 3*NA]
+                 ├── W_b1fc[AH, SH2] ──→ h_b1[B, AH] ──→ W_b1out ──→ mag_logits[B, 3*NA]
+                 ├── W_b2fc[AH, SH2] ──→ h_b2[B, AH] ──→ W_b2out ──→ ord_logits[B, 3*NA]
+                 └── W_b3fc[AH, SH2] ──→ h_b3[B, AH] ──→ W_b3out ──→ urg_logits[B, 3*NA]
+```
+
+## Proposed Forward Pass
+
+```
+h_s2[B, SH2]  ──┬── W_b0fc[AH, SH2] ──→ h_b0[B, AH] ──→ W_b0out ──→ dir_logits[B, 3*NA]
+                 │                                                          │
+                 │   ┌─── E[Q_dir][B, 3] ←── softmax(dir_logits) · z ──────┘
+                 │   │
+                 ├── [h_s2; Q_dir][B, SH2+3] ──→ W_b1fc[AH, SH2+3] ──→ h_b1 ──→ mag_logits
+                 │
+                 ├── W_b2fc[AH, SH2] ──→ h_b2[B, AH] ──→ W_b2out ──→ ord_logits[B, 3*NA]
+                 └── W_b3fc[AH, SH2] ──→ h_b3[B, AH] ──→ W_b3out ──→ urg_logits[B, 3*NA]
+```
+
+### Key change: magnitude branch 1 now depends on branch 0's output
+
+1. Direction forward runs first (already the case — branch 0)
+2. Compute `E[Q_dir]` from direction logits: 3 scalars per sample = softmax over atoms × support
+3. Concat `[h_s2; Q_dir]` into scratch buffer `[B, SH2+3]`
+4. Magnitude FC GEMM uses wider input: K = SH2+3 instead of SH2
+
+## Q_dir Computation
+
+The direction head outputs `dir_logits[B, 3, NA]`. For each direction action a ∈ {Short, Flat, Long}:
+
+```
+p_a[j] = softmax(dir_logits[b, a, :])  over atoms j=0..NA-1
+E[Q_a] = Σ_j p_a[j] * z_j             where z_j = v_min + j * delta_z
+```
+
+This produces `Q_dir[B, 3]` — three expected Q-values per sample. These are the same computation as `compute_expected_q` but for direction only (3 actions, not 12).
+
+## Implementation
+
+### CUDA kernel: `mag_concat_qdir`
+
+Small kernel that:
+1. Reads `h_s2[B, SH2]` and `dir_logits[B, 3, NA]` + per-sample support
+2. Computes E[Q] for each of the 3 direction actions (inline softmax)
+3. Writes `[h_s2; Q_dir][B, SH2+3]` to a scratch buffer
+
+Grid: `ceil(B/256)`, Block: 256. One thread per sample, loops over 3 actions × NA atoms.
+
+### Weight changes
+
+- `w_b1fc`: `[AH, SH2]` → `[AH, SH2+3]` — 3 extra columns
+- `b_b1fc`: unchanged (`[AH]`)
+- All other weights: unchanged
+
+The 3 new columns of `w_b1fc` are initialized to zero (no conditioning bias at start — the model learns to use Q_dir gradually).
+
+### param_sizes change
+
+```rust
+// [12] w_b1fc: was adv_h * shared_h2, now adv_h * (shared_h2 + 3)
+cfg.adv_h * (cfg.shared_h2 + 3),  // [12] w_b1fc (direction-conditioned)
+```
+
+Total extra parameters: `3 * adv_h` (e.g., 3 × 256 = 768 parameters). Negligible.
+
+### Forward pass changes
+
+In `batched_forward.rs`, the branch loop dispatches GEMM per branch. For branch 1 (magnitude):
+- Input pointer: `mag_concat_ptr` instead of `h_s2_ptr`
+- GEMM K dimension: `SH2+3` instead of `SH2`
+- LDB: `SH2+3`
+- A new GEMM descriptor is needed for this shape
+
+### Backward pass changes
+
+In `batched_backward.rs`, the branch FC backward for branch 1:
+
+**dW**: `dW[AH, SH2+3] += dY^T @ X_concat` — uses the wider X_concat as input
+**dX**: `dX_concat[B, SH2+3] = dY @ W_b1fc` — produces wider gradient
+
+The wider dX_concat is split:
+- `dX_concat[:, :SH2]` → accumulated into `scratch_d_h_s2` (existing flow)
+- `dX_concat[:, SH2:]` → `d_Q_dir[B, 3]` — gradient through direction Q-values
+
+`d_Q_dir` flows back through the direction head's softmax → atoms → direction logits. This adds a second gradient contribution to branch 0's logits (in addition to the existing C51 gradient). This creates a feedback loop: direction choices affect magnitude quality → magnitude gradient flows back to direction → direction learns to be more decisive.
+
+### CUDA Graph recapture
+
+The shape change (wider GEMM for branch 1) requires graph recapture. This is a one-time cost at the start of training. The recapture happens automatically because the forward/backward functions are called outside the graph on the first invocation.
+
+### Buffer allocation
+
+One new buffer: `mag_concat_buf[B, SH2+3]` f32. Allocated in `GpuDqnTrainer` constructor. Size: `batch_size * (shared_h2 + 3) * 4` bytes.
+
+One new buffer for backward: `d_mag_concat_buf[B, SH2+3]` f32. Same size as forward concat.
+
+### Serialization
+
+Existing checkpoint format stores flat parameter buffer. With the wider `w_b1fc`, the parameter buffer is slightly larger. Old checkpoints are NOT compatible (different total_params). New training runs start fresh.
+
+## Files Changed
+
+| File | Change |
+|------|--------|
+| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | `compute_param_sizes` index [12], alloc concat buffers, pass pointers |
+| `crates/ml/src/cuda_pipeline/batched_forward.rs` | Branch 1 uses concat input, new GEMM shape, add concat kernel call |
+| `crates/ml/src/cuda_pipeline/batched_backward.rs` | Branch 1 backward uses wider X, split dX, backprop d_Q_dir |
+| `crates/ml/src/cuda_pipeline/experience_kernels.cu` | `mag_concat_qdir` kernel |
+| `crates/ml/build.rs` | No change (kernel goes in experience_kernels.cu or new file) |
+
+## Risks
+
+| Risk | Mitigation |
+|------|-----------|
+| Graph recapture on first step | One-time cost, well-understood pattern |
+| Direction → magnitude dependency breaks parallel branch dispatch | Branch 1 waits for branch 0 event, branches 2+3 still parallel |
+| Old checkpoints incompatible | New runs only, no migration needed |
+| d_Q_dir backprop through softmax adds complexity | Can be simplified: treat Q_dir as detached (no gradient to direction). Test both. |
+
+## Success Criteria
+
+- `cargo check -p ml` passes
+- 19 smoke tests pass with avg_grad > 0
+- Direction-conditioned: `Long+Full` and `Flat+Full` produce different magnitude Q-values for the same market state
+- Graph capture succeeds on both RTX 3050 (local) and H100 (production)
+
+## Open Decision: Detach vs. Backprop
+
+**Option A (simpler):** Treat `Q_dir` as detached — no gradient flows back to the direction head through this path. The direction head is trained only by C51 loss on its own logits. Magnitude gets the conditioning signal but doesn't influence direction learning.
+
+**Option B (full backprop):** `d_Q_dir` flows back through softmax → direction logits. Direction head gets an additional gradient signal: "your direction choices affect magnitude quality." This creates a cooperative loop where direction learns to be more decisive because it helps magnitude.
+
+**Recommendation:** Start with Option A (detached). If val_Sharpe improves, try Option B in a follow-up. Option A is 50% less code and eliminates the softmax backward path.