From 13773375b20d2982a7ddefe8f757debcce1e5fcb Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 27 Apr 2026 18:14:22 +0200 Subject: [PATCH] feat(moe): MoeGate + MoeExpert skeletons + DQN field plumbing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Module crates/ml-dqn/src/moe.rs introduces: - MoeGate: 42→64→8 softmax gating network, zero-init for uniform start. - MoeExpert: 256→64→256 bottleneck, Xavier init per-expert seed. - Constants: MOE_NUM_EXPERTS=8, MOE_EXPERT_BOTTLENECK=64, MOE_GATE_HIDDEN=64. DQN struct gains optional moe_gate / moe_experts fields, initialized to None for now — actual wire-up into the forward graph in Phase 3. Spec: docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md §4. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml-dqn/src/dqn.rs | 12 +++++ crates/ml-dqn/src/lib.rs | 4 ++ crates/ml-dqn/src/moe.rs | 102 ++++++++++++++++++++++++++++++++++++++ docs/dqn-wire-up-audit.md | 8 +++ 4 files changed, 126 insertions(+) create mode 100644 crates/ml-dqn/src/moe.rs diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index 71e814f5e..edd8465f5 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -877,6 +877,16 @@ pub struct DQN { /// Noise sigma scale for annealed exploration (1.0 = default, decreases over training) noise_sigma_scale: f64, + + /// Mixture-of-Experts gate subnetwork: state[42] → 64 → 8 → softmax. + /// Initialized to zero so initial g(s) = uniform 1/K for every state. + /// `None` until Phase 3 wires it up. + pub moe_gate: Option, + + /// 8 expert MLPs: each is h_s1[256] → 64 → h_s2[256] bottleneck. + /// Independent Xavier init per expert for initial differentiation. + /// `None` until Phase 3 wires it up. + pub moe_experts: Option>, } impl DQN { @@ -1032,6 +1042,8 @@ impl DQN { q_value_divergence_counter: 0, training_forward_active: false, noise_sigma_scale: 1.0, + moe_gate: None, + moe_experts: None, }) } diff --git a/crates/ml-dqn/src/lib.rs b/crates/ml-dqn/src/lib.rs index 7c0dae18b..59ea83c36 100644 --- a/crates/ml-dqn/src/lib.rs +++ b/crates/ml-dqn/src/lib.rs @@ -50,6 +50,9 @@ pub mod multi_asset; // Branching DQN (Phase C+) pub mod branching; +// Mixture-of-Experts regime gating (Phase 1 skeleton) +pub mod moe; + #[cfg(test)] mod branching_composition_tests; #[cfg(test)] @@ -113,6 +116,7 @@ pub use reward::{MarketData, RewardConfig, RewardFunction, RiskMetrics}; pub use distributional::{CategoricalDistribution, DistributionalConfig, DistributionalType}; pub use distributional_dueling::{DistributionalDuelingConfig, DistributionalDuelingQNetwork}; pub use branching::{BranchingConfig, BranchingDuelingQNetwork, BranchOutput}; +pub use moe::{MoeExpert, MoeGate, MOE_EXPERT_BOTTLENECK, MOE_GATE_HIDDEN, MOE_NUM_EXPERTS}; pub use dueling::{DuelingConfig, DuelingQNetwork}; pub use rainbow_config::{RainbowAgentConfig, RainbowAgentMetrics, RainbowDQNConfig}; pub use rainbow_network::{RainbowNetwork, RainbowNetworkConfig}; diff --git a/crates/ml-dqn/src/moe.rs b/crates/ml-dqn/src/moe.rs new file mode 100644 index 000000000..e6dda58c2 --- /dev/null +++ b/crates/ml-dqn/src/moe.rs @@ -0,0 +1,102 @@ +//! Mixture-of-Experts components for the DQN policy network. +//! +//! Replaces the vestigial RegimeConditionalDQN with a learned-gate +//! K=8 expert mixture per the spec at +//! `docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md`. + +use std::sync::Arc; + +use cudarc::driver::CudaStream; +use ml_core::cuda_autograd::GpuTensor; +use ml_core::MLError; +use rand::SeedableRng; +use rand::rngs::StdRng; +use rand::Rng; +use rand_distr::Normal; + +/// K=8 experts (compile-time constant per spec §3). +pub const MOE_NUM_EXPERTS: usize = 8; + +/// Expert bottleneck width per spec §4.2. +pub const MOE_EXPERT_BOTTLENECK: usize = 64; + +/// Gate hidden width per spec §4.1. +pub const MOE_GATE_HIDDEN: usize = 64; + +/// Gate subnetwork: state[42] → 64 → 8 → softmax over experts. +/// +/// Parameters initialized to zero so g(s) = uniform 1/K initially — +/// no expert is favored by init noise; specialization emerges from data. +pub struct MoeGate { + /// Linear 1 weights [state_dim, 64], zero-init. + pub w1: GpuTensor, + /// Linear 1 bias [64], zero-init. + pub b1: GpuTensor, + /// Linear 2 weights [64, 8], zero-init. + pub w2: GpuTensor, + /// Linear 2 bias [8], zero-init. + pub b2: GpuTensor, +} + +impl MoeGate { + /// Construct with zero-initialized weights so initial gate = uniform 1/K. + pub fn new(state_dim: usize, stream: &Arc) -> Result { + let w1 = GpuTensor::zeros(&[state_dim, MOE_GATE_HIDDEN], stream)?; + let b1 = GpuTensor::zeros(&[MOE_GATE_HIDDEN], stream)?; + let w2 = GpuTensor::zeros(&[MOE_GATE_HIDDEN, MOE_NUM_EXPERTS], stream)?; + let b2 = GpuTensor::zeros(&[MOE_NUM_EXPERTS], stream)?; + Ok(Self { w1, b1, w2, b2 }) + } +} + +/// Single expert MLP: h_s1[256] → 64 → h_s2[256] with bottleneck. +/// +/// Each expert has independent Xavier init for initial differentiation. +pub struct MoeExpert { + /// Linear 1 weights [256, 64], Xavier init. + pub w1: GpuTensor, + /// Linear 1 bias [64], zero-init. + pub b1: GpuTensor, + /// Linear 2 weights [64, 256], Xavier init. + pub w2: GpuTensor, + /// Linear 2 bias [256], zero-init. + pub b2: GpuTensor, +} + +impl MoeExpert { + /// Construct with Xavier-initialized weights. + /// + /// `seed` provides per-expert distinct random init so experts have + /// non-identical starting points. + pub fn new(h_dim: usize, seed: u64, stream: &Arc) -> Result { + let w1 = init_xavier( + &[h_dim, MOE_EXPERT_BOTTLENECK], + seed, + stream, + )?; + let b1 = GpuTensor::zeros(&[MOE_EXPERT_BOTTLENECK], stream)?; + let w2 = init_xavier( + &[MOE_EXPERT_BOTTLENECK, h_dim], + seed.wrapping_add(1), + stream, + )?; + let b2 = GpuTensor::zeros(&[h_dim], stream)?; + Ok(Self { w1, b1, w2, b2 }) + } +} + +/// Xavier-normal initialization helper. +/// +/// One-time startup cost; not on the training hot-path. +fn init_xavier(shape: &[usize], seed: u64, stream: &Arc) -> Result { + let fan_in = shape.first().copied().unwrap_or(1) as f32; + let fan_out = shape.get(1).copied().unwrap_or(1) as f32; + let stddev = (2.0_f32 / (fan_in + fan_out)).sqrt(); + let normal = Normal::new(0.0_f32, stddev) + .map_err(|e| MLError::ModelError(format!("MoE Xavier dist: {e}")))?; + let numel: usize = shape.iter().product(); + let mut rng = StdRng::seed_from_u64(seed); + let host: Vec = (0..numel).map(|_| rng.sample(normal)).collect(); + GpuTensor::from_host(&host, shape.to_vec(), stream) + .map_err(|e| MLError::ModelError(format!("MoE init xavier: {e}"))) +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 611e60882..23a11e27e 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -759,3 +759,11 @@ The 3 remaining Orphan rows are: + +