From afb37b3a3101bbb64ee45e15e9158168351d0bcf Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 20 Apr 2026 07:50:44 +0200 Subject: [PATCH] =?UTF-8?q?fix:=20zero-init=20W=5FO=20and=20OFI=20embed=20?= =?UTF-8?q?weights=20=E2=80=94=20prevent=20trunk=20feature=20corruption?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Attention W_O was Xavier-initialized, injecting ~sqrt(D) magnitude noise into the residual connection. With the DtoD→copy_f32 fix making attention active (was a no-op before), random W_O corrupted all branch head inputs → MaxDD=100% on validation. Fix: W_O zero-init → attention starts as identity (output ≈ h_s2). W_Q/K/V stay Xavier (they project to SDP space, not the residual). OFI embed weights also zero-init → Mamba2 history starts as [h_s2; 0] and attention input starts as [h_s2; 0]. No noise injection at init. Both learn from zero as gradients shape them toward useful patterns. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/gpu_attention.rs | 10 +++++++--- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 17 +++++------------ 2 files changed, 12 insertions(+), 15 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_attention.rs b/crates/ml/src/cuda_pipeline/gpu_attention.rs index 5c98b9159..038892e59 100644 --- a/crates/ml/src/cuda_pipeline/gpu_attention.rs +++ b/crates/ml/src/cuda_pipeline/gpu_attention.rs @@ -329,14 +329,18 @@ impl GpuAttention { 42u64.hash(&mut hasher); let mut rng_state = hasher.finish(); - // Initialize W_Q, W_K, W_V [D, id] and W_O [D, D] with Xavier - let weight_count = 3 * d * id + d * d; - for i in 0..weight_count { + // Initialize W_Q, W_K, W_V [D, id] with Xavier — but W_O [D, D] stays ZERO. + // Zero W_O makes attention start as identity: output = LN(h_s2 + 0 @ attn) ≈ h_s2. + // This prevents random attention projections from corrupting trunk features at init. + // W_O learns from zero as attention discovers useful cross-feature patterns. + let qkv_weight_count = 3 * d * id; // W_Q + W_K + W_V only, NOT W_O + for i in 0..qkv_weight_count { rng_state = rng_state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); let u = (rng_state >> 33) as f32 / (1u64 << 31) as f32; let normal = (u - 0.5) * 3.46; host_params[i] = normal as f32 * xavier_std; } + // W_O at offset 3*D*id stays zero (from vec![0.0; total_params] init) // Biases = 0, LayerNorm gamma = 1, beta = 0 // Layout: W_Q[D*id] + b_Q[D] + W_K[D*id] + b_K[D] + W_V[D*id] + b_V[D] diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 8bcc420d9..f41ec0c9d 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -6795,18 +6795,11 @@ impl GpuDqnTrainer { // ── OFI embed MLP buffers (18→10) ── let ofi_embed_input_buf = alloc_f32(&stream, b * 18, "ofi_embed_input")?; let ofi_embed_output_buf = alloc_f32(&stream, b * 10, "ofi_embed_output")?; - let mut ofi_embed_w = alloc_f32(&stream, 10 * 18, "ofi_embed_w")?; - // Xavier init: scale = sqrt(2.0 / (fan_in + fan_out)) = sqrt(2.0 / 28) - { - let scale = (2.0_f32 / 28.0).sqrt(); - let init_data: Vec = (0..10 * 18).map(|j| { - let hash = (j as u32).wrapping_mul(2654435761).wrapping_add(0xCAFEBEEF); - let u = (hash as f32) / (u32::MAX as f32) * 2.0 - 1.0; - u * scale - }).collect(); - stream.memcpy_htod(&init_data, &mut ofi_embed_w) - .map_err(|e| MLError::ModelError(format!("ofi_embed_w xavier init: {e}")))?; - } + let ofi_embed_w = alloc_f32(&stream, 10 * 18, "ofi_embed_w")?; + // Zero-init OFI embed weights so Mamba2/attention start with clean h_s2. + // The embed produces zeros at init → h_history = [h_s2; zeros] → no corruption. + // Weights learn from zero as gradients flow back from Mamba2/attention. + // (Previous Xavier init with scale ~0.27 injected noise that corrupted trunk features.) let ofi_embed_b = alloc_f32(&stream, 10, "ofi_embed_b")?; // zero-init bias info!("GpuDqnTrainer: OFI embed MLP buffers allocated (18→10, {} params)", 10 * 18 + 10);