spec: add Components 10-12 from live H100 plateau analysis
Observed train-vpb4w freeze at epoch 23 (val_Sharpe=51.11): - avg_grad=10.0 (nonzero!) but Q-stats FROZEN for 4+ epochs - Root cause: Adam momentum converged to fixed point, not zero gradient Component 10: Targeted Adam momentum reset for branch weights only (indices 8-41). Lighter than full rewind, preserves trunk convergence. Component 11: Q-gap target attractor for spread gradient. Amplifies spread when Q-gap is below target, reduces when above. Prevents spread from being too weak at high Q-gap levels. Component 12: Cosine LR schedule with warm restarts (T=20 epochs). LR decay prevents overshoot, restart jumps break fixed points. Half-momentum at restart preserves gradient direction. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -404,14 +404,116 @@ When utilization is high (atoms well-used), `entropy_weight ≈ 0` (no interfere
|
||||
|
||||
---
|
||||
|
||||
## Component 10: Adam Momentum Reset on Plateau Detection
|
||||
|
||||
### Problem (observed in train-vpb4w epochs 23-26)
|
||||
avg_grad = 10.0 (nonzero, constant) but Q-stats are FROZEN for 4+ epochs. The gradient is active but Adam's momentum (m_buf, v_buf) has converged to a fixed point where the gradient exactly cancels the momentum update. The weights don't change despite nonzero gradient.
|
||||
|
||||
This is a DIFFERENT freeze than train-w6qfd (which had zero gradient). Here the optimizer is stuck, not the model.
|
||||
|
||||
### Solution
|
||||
When the trajectory backtracking system detects a plateau (5 consecutive frozen epochs), BEFORE attempting a full rewind, try a lightweight fix: reset Adam momentum for the frozen branch weights only.
|
||||
|
||||
```
|
||||
// In the backtracking plateau detection path:
|
||||
if plateau_detected && !route_active:
|
||||
// Try momentum reset FIRST (cheaper than full rewind)
|
||||
// Reset m_buf and v_buf for branch head weights only (indices 8-41)
|
||||
// Keep trunk weights (indices 0-7) and bottleneck (24-25) momentum intact
|
||||
zero_adam_branch_momentum()
|
||||
|
||||
// If this doesn't help after 3 epochs, THEN trigger full rewind
|
||||
```
|
||||
|
||||
### Properties
|
||||
- Zero parameters
|
||||
- Targeted reset: only branch heads + VSN + GLU gate momentum (indices 8-41), not trunk
|
||||
- Trunk momentum carries valuable gradient history — resetting it would lose convergence
|
||||
- Branch momentum at the plateau is the problem — resetting gives fresh gradient estimates
|
||||
- Lighter than full rewind (no checkpoint restore, no weight change)
|
||||
- Can be tried before each rewind attempt as a pre-filter
|
||||
|
||||
### Implementation
|
||||
- New method `reset_branch_adam_momentum()` on GpuDqnTrainer: zeros m_buf and v_buf for byte offsets corresponding to param indices 8-41 only
|
||||
- Uses `cuMemsetD8Async` on the specific offset ranges (not the full buffer)
|
||||
- Integrated into backtracking: plateau → try momentum reset → if still frozen after 3 epochs → full rewind
|
||||
|
||||
---
|
||||
|
||||
## Component 11: Spread Gradient Q-Gap Target Attractor
|
||||
|
||||
### Problem (observed in train-vpb4w)
|
||||
The spread gradient has fixed intensity: `spread_scale = inv_batch * delta_z * liquid_mod`. At Q-gap=0.365 with inv_batch=6e-5 and delta_z=0.146, spread_scale ≈ 8.7e-6. This is too weak to break the plateau — the CE gradient at the fixed point exactly balances the spread.
|
||||
|
||||
### Solution
|
||||
Scale the spread gradient by how far Q-gap is from a target. When Q-gap is below target, amplify spread. When above, reduce. This creates an attractor that prevents both collapse AND explosion:
|
||||
|
||||
```
|
||||
target_q_gap = max(0.3, 10 * eval_q_std_ema) // adaptive target
|
||||
gap_ratio = target_q_gap / max(q_gap, 1e-6) // >1 when below target, <1 when above
|
||||
spread_amplifier = clamp(gap_ratio, 0.1, 10.0) // bounded amplification
|
||||
spread_scale = inv_batch * delta_z * liquid_mod[d] * spread_amplifier
|
||||
```
|
||||
|
||||
At the train-vpb4w plateau: Q-gap=0.365, target=0.3 → gap_ratio=0.82 → amplifier=0.82 (slight reduction, Q-gap is above target). This is correct — the model has enough spread, the issue is Adam momentum not the spread gradient.
|
||||
|
||||
But if Q-gap starts declining (early sign of collapse), the amplifier ramps up BEFORE the freeze sets in.
|
||||
|
||||
### Properties
|
||||
- Zero parameters — computed from existing Q-stats
|
||||
- `spread_amplifier` is a scalar passed to c51_grad_kernel (same pattern as liquid_mod)
|
||||
- Uses pinned device-mapped memory or device buffer
|
||||
- target_q_gap is adaptive: 10× eval_q_std_ema, floored at 0.3
|
||||
|
||||
### Implementation
|
||||
- Compute `spread_amplifier` in `update_liquid_tau` (alongside per-branch modulation)
|
||||
- Add as kernel argument to c51_grad_kernel (or multiply into liquid_mod before passing)
|
||||
- Simpler: multiply into liquid_mod[d] before writing to pinned memory: `liquid_mod[d] *= spread_amplifier`
|
||||
|
||||
---
|
||||
|
||||
## Component 12: Cosine Learning Rate Schedule with Warm Restarts
|
||||
|
||||
### Problem
|
||||
Fixed learning rate + Adam momentum convergence = fixed point. The optimizer settles into a basin and can't escape. Standard deep learning solution: cosine annealing with warm restarts (SGDR).
|
||||
|
||||
### Solution
|
||||
```
|
||||
// Cosine schedule with warm restart every T epochs:
|
||||
T_restart = 20 // restart period
|
||||
t_local = epoch % T_restart // position in current cycle
|
||||
lr_mult = 0.5 * (1 + cos(pi * t_local / T_restart)) // cosine decay [1.0 → 0.0]
|
||||
lr = base_lr * max(lr_mult, 0.1) // floor at 10% of base
|
||||
|
||||
// At each restart (t_local == 0): Adam momentum is partially reset
|
||||
// m_buf *= 0.5 (half the momentum, not full zero — preserves direction)
|
||||
```
|
||||
|
||||
### Properties
|
||||
- Zero parameters
|
||||
- LR varies from base_lr to 0.1×base_lr over T=20 epochs, then restarts
|
||||
- At restart: half the Adam momentum (not full reset — keep gradient direction)
|
||||
- The restart acts as a natural plateau breaker: LR jumps → gradient pushes harder → escapes fixed point
|
||||
- train-vpb4w froze at epoch 23 — with T=20, the first restart would be at epoch 20, potentially preventing the freeze entirely
|
||||
|
||||
### Implementation
|
||||
- Compute `lr_mult` on CPU from epoch counter (cheap)
|
||||
- Pass as scalar to Adam kernel (existing `lr` parameter)
|
||||
- At restart epochs: `m_buf *= 0.5` via a simple scale kernel (or cuBLAS SSCAL)
|
||||
- ~10 lines in training_loop.rs, no CUDA kernel changes
|
||||
|
||||
---
|
||||
|
||||
## Success Criteria
|
||||
|
||||
| Metric | Current (train-vpb4w ep 18) | Target | Source |
|
||||
| Metric | Current (train-vpb4w ep 25) | 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 |
|
||||
| Q-mean drift | +0.18 (stabilized but biased) | ±0.05 (centered) | Component 8 |
|
||||
| Atom utilization | 30% (collapsed) | >50% sustained | Component 9 |
|
||||
| Adam momentum freeze | avg_grad=10 but weights frozen | No momentum stall >3 epochs | Components 10, 12 |
|
||||
| Spread gradient effectiveness | 8.7e-6 (too weak at plateau) | Adaptive to Q-gap target | Component 11 |
|
||||
| val_Sharpe sustained | 51.11 frozen at epoch 23 | >40 sustained 50+ epochs | Components 8-12 |
|
||||
| val_Sharpe plateau | Epoch 23 freeze (train-vpb4w) | 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 |
|
||||
| Branch Q-gap variance | Uniform across branches | Direction > urgency | Components 4, 7 |
|
||||
|
||||
Reference in New Issue
Block a user