diff --git a/docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md b/docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md new file mode 100644 index 000000000..08753993c --- /dev/null +++ b/docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md @@ -0,0 +1,284 @@ +# 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[42] ─── 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[42] ── 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[42] → 64 → 8` followed by softmax produces `g[B, 8]` per state. Gate has access to the same ADX/CUSUM features the legacy threshold classifier used, plus the rest of the state. +- **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: 42·64 + 64 + 64·8 + 8 = 3,272 weights + biases. +- 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. + +## 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.