gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: zero-init W_O and OFI embed weights — prevent trunk feature corruption

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-20 07:50:44 +02:00
parent d625ca28e8
commit afb37b3a31
2 changed files with 12 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
crates/ml/src/cuda_pipeline/gpu_attention.rs
View File

@@ -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]

17
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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<f32> = (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);