spec: OOS Generalization Enhancement — 10 components for capacity redirection

3-layer defense: compress capacity (AdamW, L1-sparse, gamma anneal,
regime dropout), align objectives (transaction cost curriculum, purged
walk-forward), exploit structure (epistemic-gated magnitude, branch
independence, counterfactual augmentation, temporal consistency).

4 novel techniques (G5, G6, G9, G10), 3 novel adaptations (G7, G8, G4).
Companion spec to OOS Performance Enhancement (training dynamics).
Target: OOS Sharpe > 5 sustained across multiple market regimes.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-15 20:24:06 +02:00
parent 69a01fcf5b
commit 72283d9ebd

View File

@@ -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 |