# Mixture-of-Experts replacement for `RegimeConditionalDQN` — design

**Status:** Approved through brainstorming, ready for implementation plan.

**Date:** 2026-04-27

**Pearl:** [`pearl_learned_gate_subsumes_handcoded.md`](../../../.claude/projects/-home-jgrusewski-Work-foxhunt/memory/pearl_learned_gate_subsumes_handcoded.md) — when the network already sees the heuristic's inputs, a learned gate strictly subsumes any hand-coded discretization. This pearl emerged from the present design and is the load-bearing rationale for replacing ADX/CUSUM threshold-based regime routing with a learned gating network.

## 1. Background

`RegimeConditionalDQN` (in `crates/ml-dqn/src/regime_conditional.rs`) was designed as a 3-head network with regime-routed inference: hand-coded thresholds on ADX (state index 40) and CUSUM (state index 41) classify each state as Trending / Ranging / Volatile, and `select_action`/`forward` route to the matching head.

End-to-end investigation (2026-04-27) confirmed the architecture is **vestigial decoration**:

- The 3 heads are constructed at training start with random weights.
- All experience replay inserts go to `trending_head.memory` only (`training_loop.rs:1656`, via the `primary_dqn_mut` accessor that hardcodes to trending).
- `GpuDqnTrainer` (16k+ LOC of production training kernels) has zero references to `RegimeType`, `trending_head`, `ranging_head`, `volatile_head` — the GPU trainer operates on a single flat `params_buf`.
- Only the trending head's parameters are ever touched by gradient updates. `ranging_head` and `volatile_head` stay at random init for the entire training run.
- Checkpoints save 3 files (`{path}_trending.safetensors`, `_ranging`, `_volatile`) but only the trending one contains trained weights.
- Several support APIs (`get_count_bonuses_branched`, `config()`, `get_state_dim`) hardcode-delegate to `trending_head`, ignoring the regime split entirely.

The CPU-side regime-routing scaffold is therefore policy-blind in production: training is effectively single-network, with the appearance of 3-way regime conditioning. This violates `feedback_no_hiding.md` (wire it up or delete it) and `feedback_no_functionality_removal.md` (the user's preference is to fix not delete — but "fix" requires actually wiring the feature, not preserving the appearance of it).

## 2. Decision

Replace `RegimeConditionalDQN` with a **Mixture-of-Experts (MoE)** policy network integrated into the existing single-DQN training path. The pearl says we should not gate on the hand-coded ADX/CUSUM threshold classifier when the network already sees ADX and CUSUM in the input vector — instead, a learned gating network on the full state vector strictly subsumes the threshold classifier's information content and decides per-state how to mix K=8 specialized expert MLPs.

The legacy 3-head infrastructure, per-regime checkpoint format, regime classifier (`RegimeType::classify_from_features`, `classify_regime_masks_gpu`), and threshold config fields (`regime_adx_idx`, `regime_cusum_idx`, `regime_adx_threshold`, `regime_cusum_threshold`) are deleted in the same commit — no fallback paths, no shadow infrastructure, no backward-compat for old checkpoints. Per `feedback_no_partial_refactor.md` the migration is atomic.

## 3. Architecture overview

```
state[STATE_DIM=128] ─── shared GRN trunk ─── h_s1[256] ──┬── expert_MLP_0 ──┐
                                                ├── expert_MLP_1 ──┤
                                                ├── ...            ├── (g · {expert_k}) ── h_s2[256] ── 4 branching heads ── C51 atoms ── existing IQN/CQL/MSE/Ensemble losses
                                                └── expert_MLP_7 ──┘

                                   gate: state[STATE_DIM=128] ── small_MLP(42→64→8) ── softmax ── g[8]
```

- **Shared early trunk**: existing GRN block (unchanged) computes `h_s1`.
- **K=8 expert MLPs**: each is a small bottleneck projection `h_s1[256] → 64 → h_s2[256]`. Learns a regime-conditioned policy-relevant representation.
- **Gating network**: small MLP `state[STATE_DIM=128] → 64 → 8` followed by softmax produces `g[B, 8]` per state. Gate consumes the full state vector (`ml_core::state_layout::STATE_DIM=128`, which includes 42 bar features + portfolio + MTF + OFI features padded to 128 for cuBLAS alignment). The 42-dim subset matters because that's where ADX (40) and CUSUM (41) live — the legacy threshold classifier's inputs — but the gate sees the full state including position/risk context.
- **Mixture**: `h_s2[b, :] = Σ_k g[b, k] * expert_k(h_s1[b, :])`. Soft full mixture (no top-k hardcoding); the gate's softmax distribution emerges peaky or flat from data, no architectural sparsity constraint.
- **Downstream**: existing 4 branching heads (direction × magnitude × order × urgency), C51 distributional Q, IQN dual head — all unchanged, consume `h_s2` exactly as before.
- **Load-balancing aux loss**: `λ · K · Σ_k (mean_b g[b, k])²` with default λ=0.01. **λ is a configurable hyperparameter** (lives in `DQNHyperparameters` alongside `cql_alpha`, `iqn_lambda`, etc.) — not a hardcoded kernel constant — so L40S validation can adjust it (e.g., raise to 0.05) without re-compilation if anti-collapse is too weak. Anti-collapse only — prevents init-noise-dominated single-expert lock-in (M1) without forcing uniform utilization (M2). User-confirmed empirical signal: collapses don't recover well in this codebase, so anti-collapse insurance is non-negotiable.

## 4. Component design

### 4.1 Gate subnetwork

- Input: full 42-dim state vector.
- Architecture: `Linear(42 → 64) → LeakyReLU(α=0.01) → Linear(64 → 8) → softmax`.
- Parameters: 128·64 + 64 + 64·8 + 8 = 8,776 weights + biases (using `STATE_DIM=128` for the gate input — see §3 architecture note on the 42 vs 128 distinction).
- Output: `g [B, 8]`, rows sum to 1.
- Initialization: zero weights + zero bias on both linear layers, so initial `g(s) = 1/8` uniform for every state. No expert is favored by init noise; specialization emerges from data + the load-balancing aux.

### 4.2 Expert MLPs (× 8)

- Input: `h_s1 [B, 256]` from shared trunk.
- Architecture per expert: `Linear(256 → 64) → LeakyReLU(α=0.01) → Linear(64 → 256)`.
- Parameters per expert: 256·64 + 64 + 64·256 + 256 = 33,088 weights + biases.
- Total expert parameters: 8 × 33,088 = 264,704.
- Initialization: independent Xavier per expert, distinct seeds. Initial differentiation is small but non-zero, giving the gate a meaningful gradient signal from step 1.

### 4.3 New hyperparameter: `moe_lambda`

Added to `DQNHyperparameters` alongside `cql_alpha`, `iqn_lambda`, etc. Default `0.01`. Search-space range `[0.001, 0.1]` for hyperopt. Plumbed through `from_hyperparams` into `GpuDqnTrainConfig::moe_lambda` and into the `moe_load_balance_loss` kernel call site as a per-step f32 scalar (not graph-captured as a constant — allows runtime adjustment).

### 4.4 Existing components retained unchanged

- Shared GRN trunk (computes `h_s1`).
- 4 branching heads (direction × magnitude × order × urgency).
- C51 distributional Q (51 atoms × 4 branch sizes).
- IQN dual head (`gpu_iqn_head` + `iqn_dual_head_kernel`, FIXED_TAUS schedule).
- CQL, ensemble, MSE warmup, Polyak target update, all loss machinery.
- Replay buffer (PER, single shared, no per-regime tagging).

### 4.5 Existing components deleted (no fallback)

- `crates/ml-dqn/src/regime_conditional.rs` — entire ~700 LOC module.
- `RegimeConditionalDQN`, `RegimeType` enum, `RegimeMetrics`, `RegimeClassConfig`, `classify_from_features`, `classify_regime_masks_gpu`.
- Re-exports of the above from `crates/ml-dqn/src/lib.rs`.
- `DQNConfig::regime_adx_idx`, `regime_cusum_idx`, `regime_adx_threshold`, `regime_cusum_threshold` — 4 fields, removed from struct + 4 default builders.
- `DQNAgentType` rewritten to wrap a single `DQN` directly (no `agent: RegimeConditionalDQN` field, no per-regime delegation methods, no `get_count_bonuses_branched` returning hardcoded `None`).
- Per-regime checkpoint save format (3 files) — replaced with single safetensors file.
- Tests asserting the 3-head structure or `RegimeType` enum.

### 4.6 ISV slot extension

Current `ISV_TOTAL_DIM = 118` (after `eca26a1fe`). MoE adds 9 slots:

| Index | Name | Initial / Reset | Description |
|---|---|---|---|
| 118–125 | `MOE_EXPERT_UTIL_EMA[0..8]` | 0.125 (= 1/K uniform) | Per-expert gate-weight EMA (α=0.05), GPU-driven |
| 126 | `MOE_GATE_ENTROPY_EMA` | ln(8) ≈ 2.0794 | Entropy of batch-averaged gate distribution, EMA α=0.05 |

New `ISV_TOTAL_DIM = 127`. Producer kernel `moe_expert_util_ema_update` (single-thread, single-block, cold-path cadence — same shape as existing `h_s2_rms_ema_update`) reads the captured `gate` buffer post-forward, computes batch means + entropy, EMAs into the ISV slots.

Per `pearl_cold_path_no_exception_to_gpu_drives.md` and `feedback_no_cpu_forwards.md`, all GPU-resident, CPU read-only.

Fold-boundary reset registered in `state_reset_registry.rs` per the existing pattern.

## 5. Data flow

### 5.1 Forward (per training step)

1. State batch `[B, 42]` already in GPU (sampled by replay).
2. Shared GRN trunk → `h_s1 [B, 256]` (unchanged).
3. Gate forward: 2 cuBLAS SGEMMs + LeakyReLU + softmax → `g [B, 8]`.
4. Expert forward: 8 × (cuBLAS SGEMM + LeakyReLU + cuBLAS SGEMM) → `expert_outputs [8, B, 256]`. Implemented as either batched-gemm or 8 sequential SGEMMs (cuBLAS handles both equivalently at this scale).
5. Mixture kernel `moe_mixture_forward`: `h_s2[b, :] = Σ_k g[b, k] · e_k[b, :]`. Single launch, B·256 threads.
6. Branching heads + C51 + IQN dual head — unchanged.
7. Loss aggregation: existing C51 + IQN + CQL + ensemble + MSE warmup, plus `loss += 0.01 · 8 · Σ_k (mean_b g[b, k])²` from `moe_load_balance_loss` kernel.
8. ISV producer kernel `moe_expert_util_ema_update` launches at cold-path cadence (one per training step), GPU-driven.

### 5.2 Backward (autograd through the mixture)

- `∂h_s2 → ∂e_k`: `∂e_k[b, :] = g[b, k] · ∂h_s2[b, :]`. Each expert's gradient is gated by its mixture weight — experts the gate doesn't pick get small gradient.
- `∂h_s2 → ∂g`: `∂g[b, k] = ⟨e_k[b, :], ∂h_s2[b, :]⟩`. Plus `+λ · 2K · mean_b g[b, k] / B` from the load-balancing aux loss.
- Both computed in `moe_mixture_backward` kernel.
- `∂g → ∂(gate_logits)`: standard softmax-cross-entropy backward.
- `∂(gate_logits) → ∂(gate_weights)`: cuBLAS SGEMM backward through gate's 2 linear layers.
- `∂e_k → ∂(expert_k_weights)`: cuBLAS SGEMM backward through each expert's 2 linear layers.

All operations capture into the existing CUDA Graph; no conditional execution, no straight-through estimators.

### 5.3 Optimizer step

Single Adam optimizer, single optimizer step per training step, updates trunk + experts + gate + branching heads atomically. Adam state `(m, v)` extended for the new tensors; gradient buffers extended; target params extended (Polyak EMA tracks all parameters including gate + experts).

### 5.4 Replay buffer

Single PER buffer owned by `DQN`. No per-regime tagging, no per-regime sampling. Gate handles routing in the forward pass.

### 5.5 Checkpoint format

Single safetensors file containing:

- `trunk.{w1,b1,w2,b2,vsn.*,glu.*}` — existing trunk tensor names.
- `gate.{w1,b1,w2,b2}` — new.
- `expert_{k}.{w1,b1,w2,b2}` for k ∈ [0, 8) — new (32 new tensors).
- `branching_head.{branch_k.*}` for k ∈ [0, 4) — existing names.
- Architecture hash extends to include `K=8`, expert bottleneck width (64), gate hidden width (64), and `ISV_TOTAL_DIM=127`.
- Old (pre-MoE) checkpoints fail to load with `"layout fingerprint mismatch — re-train with MoE architecture"`. No migration path.

## 6. GPU integration

### 6.1 `params_buf` extension

The MoE adds 18 new weight tensors + 18 new bias vectors appended to the existing layout (no insertion in the middle, all existing offsets preserved):

| Tensor | Shape | Count |
|---|---|---|
| `gate_w1` | `[42, 64]` | 1 |
| `gate_b1` | `[64]` | 1 |
| `gate_w2` | `[64, 8]` | 1 |
| `gate_b2` | `[8]` | 1 |
| `expert_{k}_w1` | `[256, 64]` | 8 |
| `expert_{k}_b1` | `[64]` | 8 |
| `expert_{k}_w2` | `[64, 256]` | 8 |
| `expert_{k}_b2` | `[256]` | 8 |

Total: 36 new tensors, ~268k weights + biases. Adam `m_buf` and `v_buf` extend the same way; gradient scratch buffers extend; `target_params_buf` mirrors. Layout fingerprint hash recomputes; old checkpoints fail loudly.

### 6.2 New kernels

Three small CUDA kernels added (cuBLAS handles all SGEMMs):

1. **`moe_mixture_forward`** — given `expert_outputs [8, B, 256]` and `gate [B, 8]`, produces `h_s2 [B, 256]`. Single launch, B·256 threads, each does 8 multiply-accumulates. ~50 LOC.
2. **`moe_mixture_backward`** — given `dh_s2 [B, 256]`, `gate [B, 8]`, `expert_outputs [8, B, 256]`, produces `de_k [8, B, 256]` and `dg [B, 8]`. ~80 LOC.
3. **`moe_load_balance_loss`** — computes `λ · K · Σ_k (mean_b g[b, k])²` and its backward. Block-reduce kernel for the batch mean. ~40 LOC.

Plus the ISV producer:

4. **`moe_expert_util_ema_update`** — single-thread single-block kernel reading `gate [B, 8]` post-forward, computing batch means + entropy, EMAing into `ISV[118..127)`. Same shape as existing `h_s2_rms_ema_update`, ~40 LOC.

All graph-capturable; cuBLAS calls capture; existing `softmax_forward`/`softmax_backward` kernels reused for the gate's softmax. No conditional execution, no straight-through gradients.

### 6.3 CUDA Graph integration

Forward graph extends with new nodes for: gate forward (2 SGEMMs + softmax), expert forward (8 pairs of SGEMMs + LeakyReLU), `moe_mixture_forward`, `moe_load_balance_loss`. Backward graph extends symmetrically. ISV producer launches in the cold-path cadence (one per training step). Graph recaptures cleanly on lr change / fold boundary per existing pattern.

### 6.4 Memory-transfer constraint: mapped pinned only

Per `feedback_no_htod_htoh_only_mapped_pinned.md`: every CPU↔GPU communication path in the MoE implementation uses **mapped pinned memory** (`cuMemHostAlloc(DEVICEMAP | PORTABLE)` + `cuMemHostGetDevicePointer_v2`). No `cudaMemcpy` HtoD, no Vec-to-Vec defensive copies, including in test code. The reference template is `MappedF32Buffer` / `MappedI32Buffer` from `crates/ml/src/trainers/dqn/distributional_q_tests.rs` (used by Test 0.F); the implementation plan promotes these to a shared `crates/ml/src/cuda_pipeline/mapped_pinned.rs` module so all kernel test wrappers and any future CPU→GPU staging share one implementation. CPU is strictly read-only on the production training surface — kernel-arg scalars (passed through the launch ABI by value) are the only host-originated values that flow into kernels; everything else either originates GPU-side (replay sample, ISV slot, prior kernel output) or reaches the kernel via a mapped pinned `dev_ptr`.

## 7. Cleanup cascade (atomic, single commit)

Per `feedback_no_partial_refactor.md`, every consumer of the deleted infrastructure migrates in the same commit:

- `crates/ml-dqn/src/regime_conditional.rs` — deleted.
- `crates/ml-dqn/src/lib.rs` — drop regime exports.
- `crates/ml/src/trainers/dqn/config.rs::DQNAgentType` — rewrite to wrap `DQN` directly. Per-regime delegation methods removed; `get_count_bonuses_branched` rewired to call DQN directly (returning the live fixed `[f32; 4]` arrays from the in-progress count_bonus refactor, no longer hardcoded `None`).
- `crates/ml/src/trainers/dqn/trainer/*` — every call site of regime-routing methods migrated to direct-DQN equivalents.
- `crates/ml-dqn/src/dqn.rs::DQNConfig` — drop 4 regime threshold fields, update 4 default builders.
- Per-regime checkpoint save in `crates/ml-dqn/src/regime_conditional.rs::save_checkpoint` — gone.
- Tests: any test referencing `RegimeType`, `RegimeConditionalDQN`, `trending_head`, etc. — deleted or migrated.
- `services/ml_training_service` and other downstream consumers — verified non-dependent (gRPC proto's `use_dueling` is independent; investigated as part of the prior `use_iqn` cleanup).

The use_* flag cleanup (the `[f32; N]` count_bonus refactor and `use_dueling`/`use_distributional`/`use_noisy_nets`/`use_branching`/`use_count_bonus`/`use_cvar_action_selection` flag removal) is a precondition for this MoE work — it lands as its own commit immediately before MoE. The MoE commit then removes the regime infrastructure and rewires `DQNAgentType` to wrap `DQN` directly, replacing the count_bonus wrapper's hardcoded `None` return with the live fixed-array delegation. If the use_* cleanup commit hasn't landed yet at the time MoE implementation starts, the implementation plan will land it first per `feedback_no_partial_refactor.md`.

## 8. Testing strategy (5 layers)

### Layer 1: kernel-level unit tests

`crates/ml/tests/moe_kernels_test.rs` (new file):

- `moe_mixture_forward` correctness vs CPU reference (B=4, K=8, C=256, tolerance 1e-6).
- `moe_mixture_backward` finite-differences gradient check (perturbation 1e-3, tolerance 1e-3).
- `moe_load_balance_loss` scalar correctness + gradient check.

Sub-second per test; run on local RTX 3050 Ti or in CI.

### Layer 2: integration smoke (existing harness, MoE-aware)

Extend `crates/ml/src/trainers/dqn/smoke_tests/multi_fold_convergence.rs`:

- After smoke completes, read 8 `MOE_EXPERT_UTIL_EMA` ISV slots.
- Assert no expert utilization < 2% (anti-collapse working).
- Assert no expert utilization > 80% (gate didn't snap to single expert).
- Assert `MOE_GATE_ENTROPY_EMA > 0.5 · ln(8) ≈ 1.04` (gate meaningfully spread).

Smoke runtime budget: ≤ 25 minutes on local hardware (10% slack over current ~20 min).

### Layer 3: gate-differentiation validation

After smoke completes, on ~1000 held-out states from the val fold:

- Compute per-state gate distribution entropy `H[b]`.
- Assert `mean_b H[b] < 0.9 · ln(8) ≈ 1.87` — gate is peakier than uniform on average.
- Bin states by ADX value (low/high quartiles), check that per-expert gate weights differ by `> 0.05` for at least one expert across the bins. Confirms gate uses state information meaningfully without pre-baking ADX threshold.

Failures are findings, not test infrastructure bugs. Report concrete numbers.

### Layer 4: L40S production validation

30-epoch L40S run via `./scripts/argo-train.sh dqn`, comparison to most recent non-MoE baseline (e.g., `train-l4j5q`).

**Pass criteria:**
- Final `val_loss` not worse than baseline by more than 2%.
- `MOE_EXPERT_UTIL` (averaged over last 5 epochs) shows ≥ 4 of 8 experts at > 5% utilization.
- HEALTH_DIAG emits the `aux_moe [util=[g0,...,g7] ent=...]` line per epoch with finite values.

**Stop-and-investigate criteria** (kill quickly per `feedback_kill_runs_on_anomaly_quickly.md`):
- `MOE_GATE_ENTROPY_EMA` collapses to < 0.1 within first 3 epochs → gate snap-locked → kill, debug, re-run.
- Any expert utilization < 0.1% for 5+ consecutive epochs → expert dead → kill, raise λ to 0.05, re-run.
- Training Sharpe regresses by > 10% vs baseline at epoch 5 → MoE adds overhead without benefit → kill, investigate gate/expert sizing.

### Layer 5: architecture-hash backward-incompat test

`cargo test -p ml-dqn architecture_hash_old_checkpoint_rejected`: load a pre-MoE checkpoint, assert error message contains `"layout fingerprint mismatch"`. Confirms no-fallback contract.

## 9. Success criteria

The MoE redesign is considered successful when, after L40S validation:

1. **All 5 testing layers pass.**
2. **MoE adds value or stays neutral.** `val_loss` not worse than baseline; ideally `val_loss` improves by ≥ 1% or `Sharpe` improves by ≥ 0.05.
3. **Experts differentiate.** ≥ 4 experts at > 5% utilization, mean gate entropy < 0.9·ln(8) on val states.
4. **Codebase is cleaner.** `regime_conditional.rs` and all dependent infrastructure removed; `DQNAgentType` is a thin wrapper over `DQN`; checkpoints are single-file; per `feedback_no_hiding.md` no ghost paths remain.
5. **HEALTH_DIAG instruments per-expert behavior** so future operators can see in real-time whether the gate is differentiating or collapsing.

## 10. Out of scope / future work

These are explicitly **not** part of this design and need their own brainstorm/spec:

- **Top-k routing** (sparse activation): only revisit if scaling to K ≥ 32 or expert MLPs grow >> the option-C size.
- **Per-expert action heads**: experts specialize at the trunk-late layer, not at the action-factorization layer. If empirical evidence suggests the action factorization itself should differ by regime, that's a separate design question.
- **Hierarchical MoE** (mixture of mixtures): same gate decides everything in this design. A two-level gating scheme is conceivable but adds complexity that should be motivated by data.
- **Regime-conditional CountBonus / NoisyNets sigma scheduler**: currently broadcast across all experts. If post-MoE evidence suggests these should also be per-expert, separate spec.
- **Pre-trained warm-start of one expert from existing single-network trending checkpoint**: not done in this design (clean random init for all 8). If startup convergence is slow, future commit can add warm-start.
- **CVaR action selection on the mixed C51 distribution**: noted in the in-progress count_bonus refactor as `cvar_alpha > 0` activates risk-averse argmax. That wiring is straightforward on the post-MoE mixed C51 atoms; rolling into a follow-up commit after MoE lands.

## 11. References

- [`pearl_learned_gate_subsumes_handcoded.md`](../../../.claude/projects/-home-jgrusewski-Work-foxhunt/memory/pearl_learned_gate_subsumes_handcoded.md) — the load-bearing pearl.
- [`feedback_no_partial_refactor.md`](../../../.claude/projects/-home-jgrusewski-Work-foxhunt/memory/feedback_no_partial_refactor.md) — atomic migration of all consumers in one commit.
- [`feedback_no_hiding.md`](../../../.claude/projects/-home-jgrusewski-Work-foxhunt/memory/feedback_no_hiding.md) — wire it up or delete it; no ghost paths.
- [`feedback_no_legacy_aliases.md`](../../../.claude/projects/-home-jgrusewski-Work-foxhunt/memory/feedback_no_legacy_aliases.md) — no deprecated wrappers; rename all call sites directly.
- [`feedback_kill_runs_on_anomaly_quickly.md`](../../../.claude/projects/-home-jgrusewski-Work-foxhunt/memory/feedback_kill_runs_on_anomaly_quickly.md) — kill criteria for L40S validation.
- [`pearl_cold_path_no_exception_to_gpu_drives.md`](../../../.claude/projects/-home-jgrusewski-Work-foxhunt/memory/pearl_cold_path_no_exception_to_gpu_drives.md) — ISV producer must be GPU-driven.
- [`feedback_no_cpu_forwards.md`](../../../.claude/projects/-home-jgrusewski-Work-foxhunt/memory/feedback_no_cpu_forwards.md) — CPU read-only on the inference/training surface.
- [`feedback_v7_gem_methodology.md`](../../../.claude/projects/-home-jgrusewski-Work-foxhunt/memory/feedback_v7_gem_methodology.md) — measure before declaring victory; L40S validation is the measurement layer.

---

End of design.
