diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index f4547fbaa..77a0aebf0 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -16938,3 +16938,44 @@ Phase 3c scope (remaining for full α activation) dispatch + verdict. Estimated remaining: ~20-30 hr engineering + ~37 min smoke wall-clock. + +### Phase 3c — α architectural finding + spec revision (2026-05-13) + +While implementing Phase 3c (captured-graph integration of α forward), a fundamental architectural mismatch was discovered: + +**Original spec**: α adds `W[a] * state_121[b]` to scalar Q values. Adam-trained W via dloss/dQ gradient. + +**Actual codebase**: C51 distributional Q-learning. Q is computed from a probability distribution over fixed atom positions: `E[Q_a] = Σ_n p_n * z_n`. The C51 loss is KL divergence over distributions, not MSE over scalars. + +**The math obstacle**: Adding a constant to logits is softmax-shift-invariant. Adding to scalar expected_Q happens post-loss-path. Either way, the C51 KL loss is invariant under scalar Q bias → W gradient = 0 → W never trains → α is mathematically ineffective. + +Verification: when attempting Phase 3c implementation, traced through the network's forward pass and confirmed that `mag_concat_qdir` (the kernel computing internal Q_dir for magnitude branch conditioning) and `compute_expected_q` (action selection) and `c51_loss_kernel` (training loss) all operate on atom positions and probability distributions. No path for a scalar W bias to influence the KL loss. + +**The principled fix** (revised α design in spec `docs/plans/2026-05-12-sp22-h6-phase3-alpha-beta.md` α section, 2026-05-13): + +Per-(b, a) **atom-position shift**: + +``` +z_n_effective[b, a, n] = z_n_base + W[a] * state_121[b] for all n +``` + +Shifts atom positions state-dependently. C51 Bellman projection re-projects target onto shifted positions → loss depends on `W[a] * state_121[b]` smoothly → W gradient flows from projection arithmetic. + +**Implementation cost**: substantially higher than original spec. + +- `compute_expected_q` — atom-shift threading (4 new args) +- `c51_loss_kernel` — projection uses shifted positions for dir branch +- `c51_grad_kernel` — backward dW + dstate from projection arithmetic +- `mag_concat_qdir` — internal Q_dir compute uses shifted positions +- `quantile_q_select` / `iqn_dual_head_kernel` — investigate dir-branch atom-position usage +- The Phase A `aux_to_q_dir_bias_kernel.cu` + `aux_to_q_dir_bias_backward_kernel.cu` become **architecturally redundant** (dW + dstate gradient integrate into `c51_grad_kernel`'s existing backward). Phase A kernels stay as compiled cubins (committed in `464bc5f7a`); future cleanup commit may delete them. + +**Initialization**: structural prior `W_init = [-0.5, 0.0, +0.5, 0.0]` reflects domain belief that aux conviction × position-sign biases toward aligned trades. Adam refines from there. + +**Scope adjustment**: estimated full Phase 3-final implementation grows from ~25-35 hr to ~40-60 hr engineering due to the multi-kernel atom-shift threading and backward integration through C51's projection arithmetic. + +**Decision rationale (why not Option E "fixed structural prior")**: a fixed-W version (no training) would activate α immediately for action selection but never learn state-dependent corrections. The user explicitly chose "no shortcuts" (2026-05-13 session). The atom-position-shift design is the only mathematically principled fix; the fixed-W version is a special case where Adam updates are skipped. Implementing the full version subsumes the fixed-W version as a config choice (set Adam learning rate for W to zero). + +**β + SP11 controller + A2 eval-side + telemetry**: unchanged from previous spec sections — those parts of Phase 3 were architecturally sound. Only the α design needed correction. + +Next session: execute the revised α design per the updated spec. The runbook (`docs/plans/2026-05-13-sp22-h6-phase3-alpha-beta-runbook.md`) needs corresponding revision before execution — the runbook's α tasks (B9 Steps 4-9, C1) describe the original scalar-bias implementation and need to be rewritten for atom-shift threading. diff --git a/docs/plans/2026-05-12-sp22-h6-phase3-alpha-beta.md b/docs/plans/2026-05-12-sp22-h6-phase3-alpha-beta.md index 76499d286..546aa0a0a 100644 --- a/docs/plans/2026-05-12-sp22-h6-phase3-alpha-beta.md +++ b/docs/plans/2026-05-12-sp22-h6-phase3-alpha-beta.md @@ -22,72 +22,78 @@ Combined per `pearl_no_deferrals_for_complementary_fixes` (distinct mechanisms, Both bypass the cold-encoder bottleneck on different axes (forward activation vs training signal). The pearl prescribes one atomic plan; falsification (or confirmation) of the combined intervention rules in or out both routes simultaneously. -## α — post-encoder bypass-head +## α — atom-position shift (architecturally revised 2026-05-13) -### Parameter +### Architectural correction: scalar Q bias is INCOMPATIBLE with C51 -New learnable buffer: `W_aux_to_Q_dir: [4]` f32 (one weight per direction action under the 4-3-3-3 factored space: Short, Hold, Long, Flat). +The original spec proposed `Q_dir[b, a] += W[a] * state_121[b]` as a post-encoder bypass. For non-distributional Q-learning (scalar Q per action, MSE loss against Bellman target), this works directly: W gets gradient via `dloss/dQ`. -Adam state buffers: `adam_m_W_aux: [4]`, `adam_v_W_aux: [4]` — both f32, zero-init. +This codebase uses **C51 distributional Q-learning**. Q is computed from a probability distribution over fixed atom positions: `E[Q_a] = Σ_n p_n(s, a) * z_n`. The C51 loss is KL divergence between predicted distribution and projected Bellman target. Adding a constant to logits is softmax-shift-invariant → probability distribution unchanged → loss unchanged → **W gradient = 0, never trains**. Verified mathematically and operationally during Phase 3b implementation (commit `cb80b74ce`). -Total new param storage: 12 floats (48 bytes). Negligible relative to existing network params. +### Corrected α: per-(b, a) atom-position shift -### Forward (added inside captured rollout + training forward graphs) - -Post-Q_dir computation, **before** Q_dir is consumed by downstream branches (mag/ord/urg conditioning): +The mathematically principled fix is to shift the atom POSITIONS state-dependently: ``` -for b in 0..B, a in 0..b0_size=4: - Q_dir[b, a] += W_aux_to_Q_dir[a] * state_121[b] +z_n_effective[b, a, n] = z_n_base + W[a] * state_121[b] for all atoms n ``` -Where `state_121[b] = batch_states[b * STATE_DIM + AUX_DIR_PROB_INDEX]` = the recentered Phase-2 aux signal already in state slot 121. +Where `z_n_base = v_min + n * dz` (existing per-branch support), W is a 4-element learnable bias (one per direction action), and `state_121[b] = batch_states[b * STATE_DIM + 121]` (encoder INPUT — bypasses cold encoder). -**Kernel**: new file `crates/ml/src/cuda_pipeline/aux_to_q_dir_bias_kernel.cu`. Single-pass, B threads × 4-action inner loop OR B×4 thread grid. No reductions, no atomics, no host branches → captured-graph-safe. +Effects: -**Read source for state_121**: the assembled `batch_states[B, STATE_DIM]` buffer (the encoder input, NOT the encoder output). This way α reads the **input dim** that the cold encoder is ignoring, bypassing the encoder entirely. State[121] in batch_states is the recentered value Phase 2 wrote. +- **Expected Q shift**: `E[Q'_a] = Σ_n p_n * z_n_effective = E[Q_a] + W[a] * state_121[b]` (intended bias). +- **C51 Bellman projection**: target distribution projects onto SHIFTED atom positions. Projection arithmetic is continuous in atom positions → loss depends on `W[a] * state_121[b]` smoothly → **W gradient flows**. +- **dW gradient**: `dloss/dW[a] = Σ_b state_121[b] * Σ_n dloss/dz_n_effective` (per-sample, per-atom). +- **dstate_121 gradient**: feed back into encoder column for slot 121 (option-(i) routing). -### Backward (training only, after C51/IQN compute dQ_dir) +### Initialization: structural prior -When the trainer computes the gradient w.r.t. Q_dir output (`dQ_dir: [B, 4]`), we need: +W is initialized to a domain-prior reflecting belief that aux conviction × position-sign biases toward aligned trades: -1. `dW_aux_to_Q_dir[a] = Σ_b state_121[b] * dQ_dir[b, a]` (4-element reduction over batch). -2. `dstate_121[b] += Σ_a W_aux_to_Q_dir[a] * dQ_dir[b, a]` (per-sample sum across 4 actions). +``` +W_init = [-0.5, 0.0, +0.5, 0.0] [Short, Hold, Long, Flat] +``` -**Kernel**: new file `crates/ml/src/cuda_pipeline/aux_to_q_dir_bias_backward_kernel.cu`. Two passes (or one with two grid dimensions): -- Pass A: per-action reduction `Σ_b` → 4 scalars in `dW_aux_to_Q_dir`. Block-reduce per action; B threads each fold into per-action shared-mem, then warp/block reduce. -- Pass B: per-sample `Σ_a` → B scalars added to `dstate_121` (which lives at offset 121 of the per-sample `dbatch_states` buffer, IF that backward signal is consumed by the encoder backward — see "Gradient routing to state[121]" below). +Interpretation: `state_121 ∈ [-1, +1]` is recentered aux p_up. When `state_121 = +1` (aux predicts up): +- `W[Short] * state_121 = -0.5` → Q_short atoms shift down by 0.5 → action selection avoids Short. +- `W[Long] * state_121 = +0.5` → Q_long atoms shift up by 0.5 → action selection biases toward Long. +- Hold / Flat unchanged. -Per `feedback_no_atomicadd`: Pass A uses block tree-reduce (4 actions → 4 blocks of size 256, each block reduces its action's column). Pass B is per-sample (no reduction, just multiply-add). +Adam then refines W over training. Mathematically equivalent to `W = W_prior + W_learned` with W_learned starting at zero — the prior is just the initial value of the unified W tensor. -### Gradient routing to state[121] +### Param storage -Open architectural question: does the encoder's backward path **read** `dbatch_states[121]` to update the encoder's weight for state-slot-121? +| Buffer | Shape | Purpose | +|---|---|---| +| `W_aux_to_Q_dir` | `[4]` f32 | Adam-trained weight; init `[-0.5, 0.0, +0.5, 0.0]` | +| `adam_m_W_aux` | `[4]` f32 | Adam first moment, zero-init | +| `adam_v_W_aux` | `[4]` f32 | Adam second moment, zero-init | +| `dW_aux_buf` | `[4]` f32 | Gradient accumulator, written by backward | -Three options: -- **(i) Yes — feed the gradient back through the encoder.** The α weights AND the encoder weights both update from the aux signal. The encoder might eventually learn to use state[121] independently. Risk: gradient interference; if α and encoder learn conflicting mappings, they fight. -- **(ii) No — block gradient flow into encoder via state[121].** α is the ONLY path that updates from the aux signal. The encoder's state[121] weights stay near-zero. Cleaner attribution; α experiment is purely an α test. -- **(iii) Partial — feed gradient into the encoder only for the dim 121 weight column.** Surgical; gives encoder a path without contaminating other dims. +Total: 16 floats. Negligible. -**Decision: option (i)** — let both paths get gradient. Reasons: -- Simplest implementation (just write `dbatch_states[121] +=` and let the existing backward chain do its job). -- The encoder's state[121] weight starting near-zero means its gradient contribution is small initially; α will dominate the Q_dir adjustment. As training progresses, IF the encoder learns useful weights for state[121], both paths contribute. This is the "belt + suspenders" outcome. -- Empirical check via post-smoke telemetry: compare `W_aux_to_Q_dir` magnitudes against the encoder's column-norm for state[121]. +### Kernel inventory (atom-position-shift implementation) + +Threads `aux_atom_shift[b, a] = W[a] * state_121[b]` through every kernel that uses atom positions for the direction branch: + +| Kernel | Direction-branch atom consumer? | Phase 3-final change | +|---|---|---| +| `compute_expected_q` (`experience_kernels.cu:4932`) | YES | New args: W, batch_states, aux_dir_prob_index, state_dim, b0_size. When `d == 0`, add `aux_atom_shift[b, a]` to `z_val` per atom. | +| `c51_loss_kernel` (`c51_loss_kernel.cu`) | YES — Bellman projection reads atom positions | Same atom-shift threading. Projection uses shifted positions for dir branch. | +| `c51_grad_kernel` (`c51_grad_kernel.cu`) | YES — gradients on logits via projection | Compute `dz_shift[b, a] = Σ_n dloss/dz_n_effective`. Accumulate `dW[a] += Σ_b state_121[b] * dz_shift[b, a]` and `dstate_121[b] += Σ_a W[a] * dz_shift[b, a]`. Per `feedback_no_atomicadd`: block tree-reduce. | +| `mag_concat_qdir` (`experience_kernels.cu:5505`) | YES — internal Q_dir compute uses atom positions | Apply same atom shift inline so magnitude-branch conditioning sees α-biased Q_dir consistent with action selection. | +| `quantile_q_select` / `iqn_dual_head_kernel` | Investigate during implementation | Same shift if applicable. | + +The Phase A `aux_to_q_dir_bias_kernel.cu` + `aux_to_q_dir_bias_backward_kernel.cu` are **architecturally redundant** under this revised design. The dW + dstate gradient computations integrate into `c51_grad_kernel`'s projection backward. The Phase A kernels stay as compiled cubins (committed in `464bc5f7a`, dead code at runtime); a future cleanup commit may delete them. ### Adam wireup -Add to the trainer's Adam-step kernel call site: -- Read current `W_aux_to_Q_dir`, current `dW_aux_to_Q_dir`, current `m`, `v`. -- Standard Adam update with shared β1=0.9, β2=0.999, ε=1e-8, lr=current learning rate. -- Write back updated `W`, `m`, `v`. +Same generic Adam-step kernel as other params. Add `W_aux_to_Q_dir` + `adam_m_W_aux` + `adam_v_W_aux` to the Adam param list. β1=0.9, β2=0.999, ε=1e-8, lr from learning-rate schedule. No new kernel. -The Adam-step is already a generic per-param update kernel — extend the param list to include the new 4-element buffer. No new kernel needed for Adam itself. +### Gradient routing to state[121] -### Initialization - -`W_aux_to_Q_dir = [0.0, 0.0, 0.0, 0.0]` (4 zeros). Let gradient descent learn the per-action signs and magnitudes from data. - -Rationale: starting at zero means Phase-3 forward outputs are IDENTICAL to Phase-2 outputs on step 0. Phase-3's effect grows as Adam tunes W. This makes the experiment cleanly comparable to Phase 2 baseline — any deviation in WR is attributable to W having moved off zero. +Same option (i): `dbatch_states[b * state_dim + 121] += Σ_a W[a] * dz_shift[b, a]`. Both α and encoder paths get gradient signal. ## β — event-driven aux-aligned trade-close reward