# OOS Generalization Enhancement — Design Spec

**Goal:** Achieve val/OOS Sharpe gap < 15% on ES futures across multiple market regimes by adding 15 composable generalization components that redirect existing model capacity (500K params) from noise memorization toward robust, cost-aware, Sharpe-aligned 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 layers w_s1/b_s1/w_s2/b_s2 + value head w_v1/b_v1/w_v2/b_v2), 0.0 for indices 8+ (branch heads, VSN, KAN, regime gate, atoms, etc.). Trunk + value head are shared infrastructure that benefits from compression; branch heads need full capacity for their specialized roles.
- ~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, strictly chronological):
  Fold 1: Train [1,2,3] | Purge 100 | Val [4] | Test [5,6]
  Fold 2: Train [1,2,3,4] | Purge 100 | Val [5] | Test [6]
  Fold 3: Train [2,3,4] | Purge 100 | Val [5] | Test [6]
  
Training always precedes validation chronologically. Test is always
the last window(s). No future data leaks into training.
```

### 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 G11: Q-Value Anchoring to Flat Baseline

### Problem
Q-values represent absolute expected return, including the market's natural drift. The model wastes capacity modeling common-mode drift instead of the alpha signal. Q-mean centering removes the batch mean but doesn't anchor to a meaningful baseline.

### Solution
Subtract the Flat action's Q-value from all actions after compute_expected_q. The model predicts EXCESS return over doing nothing.

```
// After compute_expected_q:
Q_flat = Q[i, flat_action_index]  // Flat = action where dir=1 (Flat)
for a in 0..total_actions:
    Q[i, a] -= Q_flat
```

### Implementation
- ~5 lines in a new kernel or inline in compute_expected_q.
- `flat_action_index`: direction=Flat(1), magnitude=Small(0), order=0, urgency=0 → index = 1*27 + 0*9 + 0*3 + 0 = 27. But with factored Q (per-branch), anchor each branch to its "neutral" action: direction anchored to Flat(1), magnitude to Small(0), order to Market(0), urgency to Normal(0).
- Runs after Q-mean centering, before regime gate.
- Novelty: action-specific baseline in factored trading DQN.

---

## Component G12: Predictive Coding Auxiliary Loss

### Problem
The trunk learns features from Q-value gradients, which are noisy (bootstrapped). A self-supervised signal would give the trunk cleaner, stabler feature learning.

### Solution
Predict h_s2(t+1) from the Mamba2 temporal context. The prediction error provides a noise-free learning signal that naturally suppresses random/non-predictable features.

```
h_s2_predicted = W_pred @ mamba2_context   // [SH2]
pred_loss = MSE(h_s2_predicted, h_s2_actual_next)
total_loss += lambda_pred * pred_loss      // lambda=0.1
```

### Implementation
- W_pred [SH2, MAMBA2_STATE_DIM] = 256 × 16 = 4096 params. Separate buffer with standalone SGD (same as Mamba2 params).
- New CUDA kernel `predictive_coding_loss` (~20 lines): reads Mamba2 state, predicts next h_s2, computes MSE.
- Prediction error doubles as a regime change detector: high error = novel state.
- Runs every training step after Mamba2 forward.
- Novelty: predictive coding on Mamba2 temporal context in trading RL.

---

## Component G13: Sharpe-Aware Reward Shaping

### Problem
The reward is `next_close - close` (raw return). The goal is Sharpe ratio. High-return high-variance strategies score well on raw return but have low Sharpe. The reward doesn't penalize variance, so the model finds strategies that look good on one window and terrible on another.

### Solution
Replace raw return with Sharpe contribution: normalize by rolling volatility.

```
reward_sharpe = (return - return_ema) / max(return_std_ema, epsilon)
```

### Implementation
- `return_ema` and `return_std_ema` are already tracked in the reward normalization pipeline (reward_shaping_kernel.cu uses `std_ema`).
- Modify the reward computation to divide by volatility BEFORE DSR shaping (not after — DSR should shape the risk-adjusted reward).
- ~15 lines in reward_shaping_kernel.cu.
- Interacts with G2 (cost curriculum): Sharpe-aware reward + cost curriculum = the model optimizes for cost-adjusted Sharpe from the start.

---

## Component G14: Confidence-Weighted Experience Replay

### Problem
PER weights samples by TD error. High TD error includes both genuine learning opportunities AND noisy outliers. The model spends training time on random noise samples early in training, seeding memorization.

### Solution
Multiply PER priority by inverse ensemble disagreement. High TD error + low disagreement = genuine opportunity. High TD error + high disagreement = noise.

```
priority = |TD_error| * (1.0 - 0.5 * ensemble_var / var_ema)
```

### Implementation
- ~10 lines in the PER priority update path (after TD error computation).
- Uses existing `ensemble_var` from the ensemble aggregate kernel and `var_ema` from G5.
- When `ensemble_count=1`, falls back to standard PER (no variance available).
- Interacts with G5 (epistemic gate): both use ensemble variance — shared infrastructure.

---

## Component G15: Action Commitment Penalty

### Problem
The model may learn to flip positions frequently even after transaction costs (G2). Frequent flipping indicates the model is uncertain about direction, not that it has a strategy. Each flip adds noise to the reward signal.

### Solution
Penalize action changes beyond the transaction cost. The penalty is proportional to position change magnitude and inversely proportional to expected holding period.

```
commitment_penalty = lambda_commit * (1 - exp(-|position_change| / tau))
reward -= commitment_penalty
// lambda_commit = 0.01, tau = 5.0 (expected holding period in bars)
```

### Implementation
- ~5 lines in the env_step kernel, after the position update and before reward accumulation.
- `tau = 5.0` bars (calibrated from the average profitable trade duration in backtests).
- Additive with G2 transaction costs — commitment penalty is a soft holding period enforcer.
- Prevents the model from learning "oscillating" strategies that look good on Sharpe but are untradeable.

---

## 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
- G11 (Q-anchoring) + G1 (AdamW) = anchored + decayed weights → focused alpha signal
- G12 (predictive coding) + Mamba2 (existing) = self-supervised temporal features + RL features → richer trunk
- G13 (Sharpe reward) + G2 (cost curriculum) = optimizes cost-adjusted Sharpe from epoch 1
- G14 (confidence replay) + G5 (epistemic gate) = shared ensemble variance infrastructure, complementary use
- G15 (commitment) + G2 (cost curriculum) = double anti-churn: cost + holding period enforcement

### 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) + G8 (L1-sparse) — 10 lines CUDA, capacity compression + feature selection
2. G13 (Sharpe reward) — 15 lines CUDA, aligns reward with Sharpe goal
3. G2 (Transaction cost curriculum) — 25 lines, realistic cost model
4. G15 (Commitment penalty) — 5 lines CUDA, anti-churn
5. G11 (Q-anchoring) — 5 lines CUDA, focus on alpha signal
6. G3 (Purged walk-forward) — 200 lines Rust, multi-regime validation
7. G4 (Gamma annealing) — 20 lines Rust, reduces horizon noise
8. G9 (Regime dropout) — 20 lines CUDA kernel
9. G7 (Counterfactual augmentation) — 15 lines CUDA
10. G5 (Epistemic magnitude gate) — 25 lines CUDA kernel
11. G6 (Branch independence) — 30 lines CUDA kernel
12. G10 (Temporal consistency) — 30 lines CUDA kernel
13. G14 (Confidence replay) — 10 lines, PER priority adjustment
14. G12 (Predictive coding) — 30 lines CUDA + 4K params

---

## Success Criteria

All 10 components deploy together as a single release.

| Metric | Current | Target |
|--------|---------|--------|
| **val/OOS Sharpe gap** | **val=45, OOS≈0 (∞ gap)** | **< 15% gap** |
| OOS Sharpe (min across windows) | ~0 | > 10.0 sustained |
| val_Sharpe (post-generalization) | 45 (overfitted) | 15-25 (genuine edge) |
| Training Sharpe stability | -1.2 to +0.9 | > 1.0 sustained |
| 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 |
| Q-anchoring removes useful common mode | Only anchors within branches, not globally — relative ordering preserved |
| Predictive coding loss dominates C51 loss | lambda_pred=0.1 is small; auxiliary loss, not primary |
| Sharpe reward division by zero | std_ema clamped to epsilon; warmup period uses raw returns |
| Confidence replay kills exploration | Only 50% weight reduction on uncertain samples, not full removal |
| Commitment penalty prevents all trading | lambda_commit=0.01 is mild; interacts with cost curriculum anneal |
