spec: add Components 8+9 — Q-mean drift correction + atom utilization recovery
Observed in train-vpb4w (live H100 run): - Q-mean drifts 0→+0.47 over 18 epochs (bootstrapping bias) - Atom utilization drops 100%→20% (distributional collapse) - val_Sharpe peaks 54.59 then decays to 31.52 Component 8: Mean-center Q-values before action selection pipeline. Zero params, preserves ranking, prevents drift accumulation. Component 9: Adaptive atom entropy regularization. Ramps entropy_coeff when utilization drops below 50%. Uses existing c51_grad_kernel param. Both are zero-param fixes moved to top of implementation order. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -324,22 +324,93 @@ trajectory_backtracking(context_divergence) [CPBI backtrack]
|
||||
|
||||
---
|
||||
|
||||
## Component 8: Q-Mean Drift Correction
|
||||
|
||||
### Problem (observed in train-vpb4w)
|
||||
Q-mean drifts 0.0 → +0.47 over 18 epochs. All Q-values shift positive due to bootstrapping bias — the DDQN target overestimates returns. This causes the model to become optimistically biased and eventually degrades action selection quality (val_Sharpe 54→31 at epoch 16-18).
|
||||
|
||||
### Solution
|
||||
Mean-centering Q-values per training step. After `compute_expected_q` produces Q_raw[B, 12]:
|
||||
|
||||
```
|
||||
Q_mean_all = mean(Q_raw) // scalar mean across all actions and samples
|
||||
Q_centered[b, a] = Q_raw[b, a] - Q_mean_all // zero-mean Q-values
|
||||
```
|
||||
|
||||
This removes the global bias while preserving relative Q-value ordering (which action is best). The centering is applied BEFORE Q-attention, graph message passing, and action selection — so all downstream components see zero-mean Q-values.
|
||||
|
||||
### Properties
|
||||
- Zero parameters
|
||||
- One reduction kernel (mean over B×12) + one element-wise subtract
|
||||
- Preserves action ranking (subtracting a constant doesn't change argmax/Boltzmann)
|
||||
- Prevents bootstrapping drift from accumulating across epochs
|
||||
- The C51 loss still uses the ORIGINAL Q-values (centering is for action selection only, not for the Bellman target)
|
||||
|
||||
### Implementation
|
||||
- New CUDA kernel `q_mean_center`: compute mean of Q_raw[B, 12], subtract from all elements. One block for reduction, ceil(B*12/256) blocks for subtraction. Two-phase kernel or fused.
|
||||
- Runs AFTER `compute_expected_q`, BEFORE `cross_branch_q_attention`
|
||||
- Q-stats monitoring reads from ORIGINAL Q_raw (not centered) to track true Q-mean drift
|
||||
|
||||
---
|
||||
|
||||
## Component 9: Atom Utilization Recovery
|
||||
|
||||
### Problem (observed in train-vpb4w)
|
||||
Atom utilization drops 100% → 20% over 18 epochs. The C51 distribution collapses onto a few atoms — most atoms carry negligible probability mass. This degrades the distributional representation: less uncertainty information for variance sizing, quantile extraction, and diffusion denoising.
|
||||
|
||||
### Solution
|
||||
Atom entropy regularization with adaptive weight. Add a small loss term that penalizes low atom entropy:
|
||||
|
||||
```
|
||||
// Per-sample, per-action atom entropy (already computed in compute_expected_q)
|
||||
H = -sum(p[j] * log(p[j])) // atom entropy
|
||||
H_max = log(num_atoms) // maximum entropy (uniform)
|
||||
utilization = H / H_max // normalized [0, 1]
|
||||
|
||||
// Entropy bonus in the C51 loss (gradient pushes toward higher entropy)
|
||||
L_entropy = -entropy_weight * mean(H) // negative = maximize entropy
|
||||
entropy_weight = base_weight * (1 - utilization_ema) // adaptive: high weight when util is low
|
||||
```
|
||||
|
||||
When utilization is high (atoms well-used), `entropy_weight ≈ 0` (no interference with C51 learning). When utilization drops below a threshold, entropy weight increases to push the distribution back toward using more atoms.
|
||||
|
||||
### Properties
|
||||
- Zero extra parameters
|
||||
- Atom entropy and utilization already computed in `compute_expected_q` (atom_stats buffer)
|
||||
- Entropy gradient already exists in `c51_grad_kernel` (the `entropy_coeff` parameter)
|
||||
- The fix is to make `entropy_coeff` ADAPTIVE based on utilization, not fixed
|
||||
- The adaptive entropy_weight is computed on CPU from the Q-stats readback (utilization value) and passed as a kernel argument
|
||||
|
||||
### Implementation
|
||||
- In `reduce_current_q_stats`: read atom_utilization from q_stats readback
|
||||
- Track `utilization_ema` with adaptive alpha (same pattern as all other EMAs)
|
||||
- Compute `entropy_weight = base_entropy_coeff * max(0, 1 - utilization_ema / 0.5)` — ramps up when utilization drops below 50%, zero above 50%
|
||||
- Pass adaptive `entropy_weight` to `c51_grad_kernel` instead of the fixed `entropy_coeff` hyperparameter
|
||||
- No new CUDA kernels — uses existing `entropy_coeff` parameter in c51_grad_kernel
|
||||
|
||||
---
|
||||
|
||||
## Implementation Order
|
||||
|
||||
1. Component 6 (Quantile Q-select) — zero params, uses existing atoms, highest OOS impact
|
||||
2. Component 4 (Graph message passing) — 60 params, one kernel, structural coordination
|
||||
3. Component 1 (KAN spline gates) — replaces sigmoid in existing GLU, localized change
|
||||
4. Component 3 (Diffusion refinement) — 2,772 params, separate Adam, post-Q-pipeline
|
||||
5. Component 7 (TLOB injection) — widens 4 weight tensors, same pattern as Layer 3
|
||||
6. Component 2 (xLSTM context) — CPU-only, feeds into Components 5 and existing systems
|
||||
7. Component 5 (RK4 ODE) — depends on Component 2, pure algorithm change
|
||||
1. **Component 8 (Q-mean centering)** — zero params, one kernel, directly fixes the Q-mean drift
|
||||
2. **Component 9 (Atom entropy recovery)** — zero params, adaptive existing param, directly fixes atom collapse
|
||||
3. Component 6 (Quantile Q-select) — zero params, uses existing atoms, highest OOS impact
|
||||
4. Component 4 (Graph message passing) — 60 params, one kernel, structural coordination
|
||||
5. Component 1 (KAN spline gates) — replaces sigmoid in existing GLU, localized change
|
||||
6. Component 3 (Diffusion refinement) — 2,772 params, separate Adam, post-Q-pipeline
|
||||
7. Component 7 (TLOB injection) — widens 4 weight tensors, same pattern as Layer 3
|
||||
8. Component 2 (xLSTM context) — GPU kernels, feeds into Components 5 and existing systems
|
||||
9. Component 5 (RK4 ODE) — depends on Component 2, GPU kernel
|
||||
|
||||
---
|
||||
|
||||
## Success Criteria
|
||||
|
||||
| Metric | Current (CPBI baseline) | Target | Source |
|
||||
|--------|------------------------|--------|--------|
|
||||
| Metric | Current (train-vpb4w ep 18) | Target | Source |
|
||||
|--------|---------------------------|--------|--------|
|
||||
| Q-mean drift | +0.47 (unbounded growth) | ±0.05 (centered) | Component 8 |
|
||||
| Atom utilization | 20% (collapsing) | >50% sustained | Component 9 |
|
||||
| val_Sharpe sustained | 54 peak → 31 decay | >40 sustained 50+ epochs | Components 8, 9 |
|
||||
| val_Sharpe plateau | Epoch 22 freeze (train-w6qfd) | No freeze through 200 epochs | All components |
|
||||
| OOS Sharpe | ~0 oscillating | >0 sustained 20+ epochs | Components 1, 3, 6 |
|
||||
| Exploration quality | Epsilon-greedy / Boltzmann | Distributional (no epsilon) | Component 6 |
|
||||
@@ -353,6 +424,8 @@ trajectory_backtracking(context_divergence) [CPBI backtrack]
|
||||
|
||||
| Risk | Mitigation |
|
||||
|------|-----------|
|
||||
| Q-mean centering removes useful signal | Centering preserves relative ordering; C51 loss uses original Q-values |
|
||||
| Adaptive entropy too aggressive | Ramps linearly, zero above 50% utilization — no interference when atoms are healthy |
|
||||
| KAN spline overflow | Clamp gate output to [0, 1] |
|
||||
| xLSTM C matrix divergence | Forget gate f_t ∈ (0,1) prevents unbounded growth |
|
||||
| Diffusion refinement distorts Q ordering | Small step (alpha=0.1) + residual preserves ranking |
|
||||
|
||||
Reference in New Issue
Block a user