diff --git a/docs/superpowers/specs/2026-04-15-oos-generalization-enhancement-design.md b/docs/superpowers/specs/2026-04-15-oos-generalization-enhancement-design.md new file mode 100644 index 000000000..676ecafcd --- /dev/null +++ b/docs/superpowers/specs/2026-04-15-oos-generalization-enhancement-design.md @@ -0,0 +1,329 @@ +# OOS Generalization Enhancement — Design Spec + +**Goal:** Achieve sustained OOS Sharpe > 5 on ES futures across multiple market regimes by adding 10 composable generalization components that redirect existing model capacity (500K params) from noise memorization toward robust, cost-aware strategies. + +**Problem:** The model achieves val_Sharpe=45 on a fixed validation window but training Sharpe oscillates -1.2 to +0.9. Five H100 runs confirm: the model memorizes temporal patterns instead of learning generalizable strategies. Root causes: (1) no weight decay — every parameter free to encode noise, (2) transaction costs exist in env_step but no curriculum — model finds churn strategies early and gets stuck, (3) single validation window — checkpoint selection optimizes for overfitting, (4) gamma=0.99 injects 100-bar horizon noise, (5) no data augmentation — model memorizes temporal position, (6) no structural regularization — branches co-adapt, trunk over-specializes. + +**Architecture:** 3-layer defense composing through the existing fused CUDA training pipeline. Layer 1 compresses capacity (G1 AdamW, G4 gamma anneal, G8 L1-sparse input, G9 regime dropout). Layer 2 aligns objectives (G2 transaction cost curriculum, G3 purged walk-forward). Layer 3 exploits structure (G5 epistemic-gated magnitude, G6 branch independence, G7 counterfactual augmentation, G10 temporal consistency). + +**Companion to:** `2026-04-15-oos-performance-enhancement-design.md` (training dynamics). This spec addresses generalization. Both deploy together — the performance plan makes the model train better, this plan makes it learn the right things. + +--- + +## Component G1: AdamW Decoupled Weight Decay + +### Problem +Standard Adam with no weight decay. 500K parameters free to encode noise. The model memorizes training data patterns that don't persist OOS. + +### Solution +Replace standard Adam with AdamW (decoupled weight decay) in `dqn_utility_kernels.cu`. Apply to trunk weights only (tensor indices 0-7) — branch heads need full capacity for specialization. + +``` +// Current: theta -= lr * m_hat / (sqrt(v_hat) + eps) +// AdamW: theta -= lr * (m_hat / (sqrt(v_hat) + eps) + lambda * theta) +``` + +### Implementation +- Modify Adam kernel in `dqn_utility_kernels.cu`: add `+ lambda * theta` term gated by a `weight_decay_mask` pointer (1.0 for trunk, 0.0 for branches). +- `lambda = 1e-4`, pinned device-mapped for graph safety. Ramp from 1e-5 to 1e-4 over first 10 epochs. +- `weight_decay_mask`: [total_params] buffer, 1.0 for indices 0-7 (trunk), 0.0 for indices 8+ (branches, regime gate, atoms, etc.). +- ~5 lines CUDA change in existing kernel. + +--- + +## Component G2: Transaction Cost Curriculum + +### Problem +The env_step kernel already models `spread_cost`, but the model discovers churn strategies early (when it hasn't learned to overcome costs) and gets stuck in a "trade frequently with tiny edge" local optimum that doesn't survive real spreads. + +### Solution +Annealed cost curriculum: start with 0% of real transaction costs, ramp to 100% by epoch 20. Novel twist: tie ramp rate to ensemble disagreement — high disagreement slows the ramp. + +``` +cost_anneal = min(1.0, epoch / 20) +effective_anneal = cost_anneal * (1.0 - 0.5 * clamp(ensemble_disagreement, 0, 1)) +// env_step: all cost terms *= effective_anneal +``` + +### Implementation +- `cost_anneal_pinned: *mut f32, cost_anneal_dev_ptr: u64` — pinned device-mapped scalar. +- In env_step kernel: multiply `spread_cost`, `tx_cost_multiplier` by `cost_anneal_ptr[0]`. +- Training loop: update `cost_anneal_pinned` at epoch start based on epoch count and ensemble disagreement. +- ~10 lines CUDA (multiply existing costs by factor), ~15 lines Rust (anneal schedule). + +### No-Trade Collapse Mitigation +If cost ramp causes the model to learn "do nothing": +- Ensemble disagreement automatically slows the anneal (high uncertainty = model not ready for real costs). +- Existing entropy bonus in Expected SARSA softmax maintains exploration pressure. +- Existing adaptive DSR (Component 2 from performance plan) increases DSR when Sharpe goes negative, forcing activity. + +--- + +## Component G3: Purged Walk-Forward Validation + +### Problem +Single fixed validation window. val_Sharpe=45 measures fit to ONE regime, not generalization. Checkpoint selection optimizes for overfitting. + +### Solution +Chronological purged walk-forward: split 4M bars into K=6 non-overlapping windows. Train on windows 1..K-2, purge 100 bars, validate on K-1, test on K. Rotate K times. Report min(Sharpe) across all test windows. + +``` +Window layout (K=6, each ~667K bars): + Fold 1: Train [1,2,3,4] | Purge 100 | Val [5] | Test [6] + Fold 2: Train [1,2,3,5] | Purge 100 | Val [4] | Test [6] + ...etc (always test on last window for forward-looking honesty) +``` + +### Implementation +- Modify validation split in training_loop.rs: replace single window with K-fold chronological split. +- 100-bar purge gap between train and validation sets prevents Mamba2 temporal state leakage. +- Checkpoint selection metric: `min(Sharpe across K validation windows)` replaces `val_Sharpe`. +- Backtracking system uses min-Sharpe for plateau detection: plateau when min-Sharpe frozen across 5+ epochs. +- ~200 lines Rust (validation split logic + min-Sharpe aggregation). Zero CUDA changes. +- Deflated Sharpe Ratio: penalize low trade count. `adjusted_sharpe = sharpe * sqrt(num_trades / expected_trades)`. + +--- + +## Component G4: Gamma Annealing with Auto C51 Rescaling + +### Problem +gamma=0.99 gives effective horizon ~100 bars. For mid-frequency ES trading, this injects enormous noise from unpredictable far-future returns. Q-values become dominated by long-horizon uncertainty. + +### Solution +Adaptive gamma: start at 0.90 (10-bar horizon), extend toward 0.95 as atom utilization stays healthy. Auto-rescale C51 v_min/v_max when gamma changes. + +``` +if atom_utilization_ema > 0.6: + gamma = min(0.95, gamma + 0.005) // slow extension +elif atom_utilization_ema < 0.4: + gamma = max(0.90, gamma - 0.01) // fast contraction +// else: hold current gamma (hysteresis band [0.4, 0.6]) + +// Auto-rescale C51 bounds: +R_max = max_single_step_return // from reward statistics +v_max = R_max / (1 - gamma) +v_min = -v_max +``` + +### Implementation +- New fields: `adaptive_gamma: f32`, `gamma_v_max: f32` in DQNTrainer. +- Update at epoch end after atom utilization is computed. +- Write new v_min/v_max to `per_sample_support` buffer (already pinned device-mapped). +- Hysteresis prevents oscillation: different thresholds for extend (0.6) and contract (0.4). +- ~20 lines Rust (gamma update + v_min/v_max rescale). C51 kernels already read v_min/v_max dynamically. + +--- + +## Component G5: Epistemic-Gated Magnitude (Novel) + +### Problem +Ensemble variance measures epistemic uncertainty (model hasn't seen this state) but is only used for exploration. The model can take Full(1.00) magnitude positions on completely unfamiliar states. + +### Solution +Ensemble disagreement directly scales magnitude branch Q-values. High epistemic uncertainty → forced toward Small(0.25). + +``` +// After ensemble aggregate computes ensemble_var[B, total_actions]: +var_mean = mean(ensemble_var[i, :]) // per-sample mean variance +threshold = var_ema // adaptive: EMA of training variance +mag_scale = sigmoid(-5.0 * (var_mean - threshold)) +// var << threshold → scale ≈ 1.0 (full magnitude freedom) +// var >> threshold → scale ≈ 0.0 (crush toward Small) + +Q_mag[a] = Q_mag[a] * mag_scale + Q_small_bias * (1 - mag_scale) +``` + +### Implementation +- New CUDA kernel `epistemic_gate_magnitude` (~25 lines): reads ensemble_var, computes per-sample scale, applies to magnitude branch Q-values. +- `var_ema` pinned device-mapped scalar, updated at epoch end. +- Runs after ensemble aggregate, before action selection. +- Only active when `ensemble_count > 1` (no-op for single-head training). +- Novelty: ensemble variance for position sizing in factored DQN — no prior work. + +--- + +## Component G6: Branch Independence Penalty (Novel) + +### Problem +4 branches share a trunk and co-adapt — direction and magnitude learn correlated features, wasting capacity on redundant information. + +### Solution +Cosine similarity penalty between branch hidden representations. Forces branches to learn orthogonal features. + +``` +independence_loss = 0 +for (i, j) in [(0,1), (0,2), (0,3), (1,2), (1,3), (2,3)]: + cos_sim = dot(h_bi, h_bj) / (|h_bi| * |h_bj| + eps) + independence_loss += cos_sim^2 + +total_loss += lambda_indep * independence_loss // lambda_indep = 0.01 +``` + +### Implementation +- New CUDA kernel `branch_independence_penalty` (~30 lines): reads 4 branch hidden buffers (save_h_b0..save_h_b3), computes 6 pairwise cosine similarities, writes scalar penalty. +- Penalty accumulated into `total_loss_dev_ptr` (same as C51 loss). +- Gradient: d(cos_sim^2)/d(h_bi) flows back through the branch forward pass. Requires saving branch activations (already done for backward). +- Novelty: structural regularization via factored independence, unique to multi-branch DQN. + +--- + +## Component G7: Counterfactual Return Augmentation (Novel Adaptation) + +### Problem +Training data is fixed. Model memorizes temporal position. No augmentation = overfitting. + +### Solution +50% probability during experience collection: flip episode returns. Negate price deltas, negate rewards, swap Long↔Short actions. Only direction branch flips — magnitude/order/urgency preserved. + +``` +if philox_uniform(sample_id, seed) < 0.5: + features[price_delta_indices] *= -1.0 + reward *= -1.0 + action_dir = 2 - action_dir // Long(2)↔Short(0), Flat(1) stays + // magnitude, order, urgency unchanged +``` + +### Implementation +- ~15 lines in existing env_step kernel (conditional flip before reward computation). +- `price_delta_indices`: features that are directional (returns, price deltas, momentum, CUSUM) get negated. Non-directional features (volume, ADX, volatility, spread, OFI magnitude) stay unchanged. The exact index list is determined by reading the feature vector layout from the data pipeline — approximately indices 0-15 (price-derived) are flipped, 16-41 (microstructure/volume/volatility) are preserved. +- Philox seed uses `sample_id + epoch * 1000000` for reproducibility. +- No new buffers or kernel handles. +- Novelty: domain randomization from robotics adapted to factored trading RL with branch-aware symmetry. + +--- + +## Component G8: L1-Sparse First Layer (Novel Combination) + +### Problem +42 features fed to the trunk. Many are noise. The model fits to all of them, including spurious correlations that don't persist OOS. + +### Solution +L1 proximal gradient on w_s1 (tensor index 0) only. Induces automatic feature selection — zeros out columns of the first weight matrix corresponding to useless features. + +``` +// After standard AdamW update for w_s1: +sign = (theta > 0) ? 1.0 : -1.0 +theta = sign * max(0, abs(theta) - lr * lambda_l1) // soft thresholding +``` + +### Implementation +- ~5 lines in Adam kernel: conditional on parameter offset being within w_s1 range. +- `lambda_l1 = 1e-3`. Applied only to w_s1 (tensor 0, size = shared_h1 × s1_input_dim). +- `l1_end_offset`: byte offset marking end of w_s1 in params_buf. +- After epoch 10, log non-zero feature count from w_s1 column norms (free interpretability). +- Combined with AdamW (G1): sparse input → compressed trunk → specialized branches. +- Novelty: L1 input + AdamW trunk for distributional DQN. Creates differentiable feature selector. + +--- + +## Component G9: Regime-Aware Dropout (Novel) + +### Problem +No dropout = no regularization on h_s2. Standard random dropout would work but wastes the regime information already in the state vector. + +### Solution +Dropout mask conditioned on market regime (ADX/CUSUM). Same regime always drops the same neurons. Forces each neuron to be useful across multiple regimes. + +``` +regime_hash = hash(quantize(ADX, 4) * 7 + quantize(CUSUM, 4)) +for j in 0..SH2: + mask = philox_uniform(regime_hash, j, epoch_seed) > drop_rate + h_s2[i, j] *= mask / (1 - drop_rate) // inverted dropout +``` + +### Implementation +- New CUDA kernel `regime_dropout` (~20 lines). +- `drop_rate = 0.15` (lighter than standard — regime conditioning provides implicit regularization). +- `quantize(x, 4)`: maps continuous feature into 4 buckets (quartiles from training data). +- Runs AFTER Mamba2 enrichment (don't disrupt temporal state), BEFORE branch heads. +- Disabled during validation/inference (standard dropout behavior). +- `epoch_seed` changes per epoch so the mask evolves — no single neuron is permanently dead for a regime. +- Novelty: regime-conditioned dropout, no prior work in any domain. + +--- + +## Component G10: Temporal Consistency Regularization (Novel) + +### Problem +Consecutive bars with small price change should produce similar Q-values. But the model can memorize specific bar positions — Q(bar_1000, a) ≠ Q(bar_1001, a) even when the states are nearly identical. + +### Solution +Penalize Q-value differences between consecutive similar states. Lipschitz-style regularizer adapted from robotics sim-to-real. + +``` +for consecutive samples i, i+1 in batch: + state_sim = cosine_similarity(h_s2[i], h_s2[i+1]) + if state_sim > 0.95: + q_diff = mean_a |Q[i, a] - Q[i+1, a]| + penalty += lambda_tc * q_diff * (state_sim - 0.95) / 0.05 +``` + +### Implementation +- New CUDA kernel `temporal_consistency_penalty` (~30 lines). +- Reads `save_h_s2` and `q_out_buf` for consecutive batch samples. +- Only fires when consecutive states are highly similar (common in financial data — most bars are small moves). +- `lambda_tc = 0.005`. Penalty strength proportional to similarity excess above 0.95 threshold. +- Penalty accumulated into loss; gradient flows through Q-values to branches/trunk. +- Novelty: Lipschitz regularization for C51 distributional trading, exploiting temporal structure of financial data. + +--- + +## Component Interactions + +### Synergistic +- G1 (AdamW) + G8 (L1-sparse) = learned input bottleneck: 42 → ~15 features → compressed trunk +- G2 (cost curriculum) + G5 (epistemic gate) = safe cost ramp: disagreement controls both anneal AND position size +- G3 (walk-forward) + G6 (branch independence) = regime-robust branches: walk-forward punishes regime-specific strategies, independence prevents co-adaptation +- G7 (counterfactual) + G9 (regime dropout) = double augmentation: 2× data via symmetry + regime-diverse representations +- G4 (gamma anneal) + G10 (temporal consistency) = horizon-aware smoothness + +### Compensating (adjust existing components) +- G1 (AdamW) + CVaR (existing) = double regularization → reduce CVaR alpha range from [0.1, 0.5] to [0.3, 0.5] +- G9 (regime dropout) + ensemble heads = redundant uncertainty → keep ensemble for G5 magnitude gating, not exploration + +--- + +## Implementation Order + +1. G1 (AdamW) — 5 lines CUDA, immediate impact on capacity compression +2. G2 (Transaction cost curriculum) — 25 lines, aligns reward with reality +3. G3 (Purged walk-forward) — 200 lines Rust, changes what we measure +4. G4 (Gamma annealing) — 20 lines Rust, reduces horizon noise +5. G8 (L1-sparse) — 5 lines CUDA, automatic feature selection +6. G9 (Regime dropout) — 20 lines CUDA kernel +7. G5 (Epistemic magnitude gate) — 25 lines CUDA kernel +8. G6 (Branch independence) — 30 lines CUDA kernel +9. G7 (Counterfactual augmentation) — 15 lines CUDA +10. G10 (Temporal consistency) — 30 lines CUDA kernel + +--- + +## Success Criteria + +All 10 components deploy together as a single release. + +| Metric | Current | Target | +|--------|---------|--------| +| OOS Sharpe (min across windows) | ~0 | > 5.0 sustained | +| Training Sharpe stability | -1.2 to +0.9 | > 1.0 sustained | +| val/OOS gap | val=45, OOS≈0 | gap < 2× | +| Active features (w_s1 non-zero columns) | 42/42 | < 20/42 | +| Branch correlation (max pairwise cos_sim) | unmeasured | < 0.1 | +| Atom utilization | 20-68% | > 70% | + +--- + +## Risks and Mitigations + +| Risk | Mitigation | +|------|-----------| +| AdamW too aggressive → Q-values collapse | Start lambda=1e-5, ramp to 1e-4 over 10 epochs | +| Cost curriculum → no-trade collapse | Ensemble disagreement slows anneal; entropy bonus keeps exploration | +| L1 kills all features → blank input | lambda_l1=1e-3 is gentle; monitor non-zero count, reduce if < 10 | +| Regime dropout destabilizes Mamba2 | Apply AFTER Mamba2 enrichment, not before | +| Counterfactual flip breaks asymmetric patterns | Only flip direction; magnitude/order/urgency preserved | +| Temporal consistency over-smooths Q | lambda_tc=0.005 is small; only fires when state_sim > 0.95 | +| Gamma annealing oscillates | Hysteresis: extend at util > 0.6, contract at util < 0.4 | +| Walk-forward reduces training data per fold | 4M bars ÷ 6 = 667K per window, still large for DQN | +| Branch independence penalty too strong → uniform branches | lambda_indep=0.01 is mild; monitor Q-gap per branch | +| Epistemic gate prevents all large trades | Adaptive threshold (var_ema) means gate only fires on unusual states |