crates/training/rtx-transformers/src/layers/mamba.rs

@@ -287,40 +287,7 @@ pub struct MambaBlock {
 impl MambaBlock {
     /// Create new `MambaBlock`
     pub fn new(config: MambaConfig, device: &Device) -> Result<Self> {
-        let d_inner = config.get_d_inner();
-        let dt_rank = config.get_dt_rank();
-
-        // Initialize parameters
-        let in_proj = Tensor::randn(&[config.d_model, d_inner * 2], device)?;
-        let conv1d_weight = Tensor::randn(&[d_inner, 1, config.d_conv], device)?;
-        let conv1d_bias = if config.conv_bias {
-            Some(Tensor::randn(&[d_inner], device)?)
-        } else {
-            None
-        };
-        let A_log = Tensor::randn(&[d_inner, config.d_state], device)?;
-        let x_proj = Tensor::randn(&[d_inner, dt_rank + 2 * config.d_state], device)?;
-        let dt_proj = Tensor::randn(&[dt_rank, d_inner], device)?;
-        let dt_bias = Tensor::randn(&[d_inner], device)?;
-        let d_skip = Tensor::randn(&[d_inner], device)?;
-        let out_proj = Tensor::randn(&[d_inner, config.d_model], device)?;
-
-        let selective_scan = SelectiveScan::new(d_inner, config.d_state);
-
-        Ok(Self {
-            config,
-            device: device.clone(),
-            in_proj,
-            conv1d_weight,
-            conv1d_bias,
-            A_log,
-            x_proj,
-            dt_proj,
-            dt_bias,
-            d_skip,
-            out_proj,
-            selective_scan,
-        })
+        Self::build(config, device, None)
     }
 
     /// Create `MambaBlock` with a deterministic seed.
@@ -335,38 +302,70 @@ impl MambaBlock {
     /// same operator-supplied `seed` produces a stable mapping across
     /// the six (or seven, with conv_bias) internal tensors.
     pub fn new_seeded(config: MambaConfig, device: &Device, seed: u64) -> Result<Self> {
-        let d_inner = config.get_d_inner();
-        let dt_rank = config.get_dt_rank();
+        Self::build(config, device, Some(seed))
+    }
 
-        // Six (or seven) per-tensor seeds derived from the operator
-        // seed. Using SplitMix64 so adjacent operator seeds don't
-        // produce correlated per-tensor seeds.
-        let mut state = seed;
-        let mut next_seed = || {
+    /// Shared constructor for [`Self::new`] (random) and
+    /// [`Self::new_seeded`] (`Some(seed)` ⇒ deterministic). Uses
+    /// canonical, numerically-stable selective-SSM initialization:
+    /// `A_log = ln(1..=d_state)` so `A = -(1..=d_state)` is bounded;
+    /// `dt_bias` set so `softplus(dt_bias) ≈ 0.01` (small, stable Δ);
+    /// `D = 1`; zero conv bias; projection weights scaled by `1/√fan_in`
+    /// (capped at 0.5). Plain `randn(0,1)` would overflow the real
+    /// `exp(Δ·A)` scan to NaN — the stub didn't care because it
+    /// discarded these weights.
+    fn build(config: MambaConfig, device: &Device, seed: Option<u64>) -> Result<Self> {
+        let d = config.get_d_inner();
+        let dt_rank = config.get_dt_rank();
+        let n = config.d_state;
+        let dbc = dt_rank + 2 * n;
+
+        let mut state = seed.unwrap_or(0xD1CE_F00D);
+        let mut rand_t = |shape: &[usize]| -> std::result::Result<Tensor, rtx_tensor::TensorError> {
             state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
             let mut z = state;
             z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
             z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
-            z ^ (z >> 31)
+            let s = z ^ (z >> 31);
+            match seed {
+                Some(_) => Tensor::randn_seeded(shape, device, s),
+                None => Tensor::randn(shape, device),
+            }
         };
+        let scaled =
+            |t: Tensor, f: f32| -> std::result::Result<Tensor, rtx_tensor::TensorError> {
+                let mut v = t.to_vec()?;
+                for x in v.iter_mut() {
+                    *x *= f;
+                }
+                Tensor::from_vec(v, t.shape().dims(), device)
+            };
+        let lin = |fan: usize| (1.0 / (fan.max(1) as f32).sqrt()).min(0.5);
 
-        let in_proj = Tensor::randn_seeded(&[config.d_model, d_inner * 2], device, next_seed())?;
-        let conv1d_weight =
-            Tensor::randn_seeded(&[d_inner, 1, config.d_conv], device, next_seed())?;
+        let in_proj = scaled(rand_t(&[config.d_model, 2 * d])?, lin(config.d_model))?;
+        let conv1d_weight = scaled(rand_t(&[d, 1, config.d_conv])?, 0.2)?;
         let conv1d_bias = if config.conv_bias {
-            Some(Tensor::randn_seeded(&[d_inner], device, next_seed())?)
+            Some(Tensor::from_vec(vec![0.0f32; d], &[d], device)?)
         } else {
             None
         };
-        let A_log = Tensor::randn_seeded(&[d_inner, config.d_state], device, next_seed())?;
-        let x_proj =
-            Tensor::randn_seeded(&[d_inner, dt_rank + 2 * config.d_state], device, next_seed())?;
-        let dt_proj = Tensor::randn_seeded(&[dt_rank, d_inner], device, next_seed())?;
-        let dt_bias = Tensor::randn_seeded(&[d_inner], device, next_seed())?;
-        let d_skip = Tensor::randn_seeded(&[d_inner], device, next_seed())?;
-        let out_proj = Tensor::randn_seeded(&[d_inner, config.d_model], device, next_seed())?;
+        // A = -exp(A_log) = -(1..=d_state) per channel (S4D-real init).
+        let mut a_v = vec![0.0f32; d * n];
+        for j in 0..d {
+            for nn in 0..n {
+                a_v[j * n + nn] = ((nn + 1) as f32).ln();
+            }
+        }
+        let a_log = Tensor::from_vec(a_v, &[d, n], device)?;
+        let x_proj = scaled(rand_t(&[d, dbc])?, lin(d))?;
+        let dt_proj = scaled(rand_t(&[dt_rank, d])?, lin(dt_rank))?;
+        // softplus(dt_bias) ≈ 0.01 ⇒ small, stable time-step.
+        let dt_bias_val = (0.01f32.exp() - 1.0).ln();
+        let dt_bias = Tensor::from_vec(vec![dt_bias_val; d], &[d], device)?;
+        let d_skip = Tensor::from_vec(vec![1.0f32; d], &[d], device)?;
+        let out_proj = scaled(rand_t(&[d, config.d_model])?, lin(d))?;
 
-        let selective_scan = SelectiveScan::new(d_inner, config.d_state);
+        let selective_scan = SelectiveScan::new(d, config.d_state);
 
         Ok(Self {
             config,
@@ -374,7 +373,7 @@ impl MambaBlock {
             in_proj,
             conv1d_weight,
             conv1d_bias,
-            A_log,
+            A_log: a_log,
             x_proj,
             dt_proj,
             dt_bias,

crates/training/rtx-transformers/tests/real_selective_scan.rs

@@ -52,10 +52,30 @@ fn perturbed_map(block: &MambaBlock, target: &str, delta: f32, dev: &Device) ->
     map
 }
 
+/// A seeded block with `dt_bias` overridden to 0 ⇒ `Δ = softplus(0) ≈
+/// 0.69`. The default init uses a deliberately small `Δ≈0.01` (the scan
+/// is near-identity — correct for training stability), which makes the
+/// scan-only weights' influence vanish below the f32 floor. This widens
+/// `Δ` so the scan genuinely contributes to the output, letting us
+/// observe each weight's effect.
+fn active_block(cfg: &MambaConfig, dev: &Device, seed: u64) -> MambaBlock {
+    let base = MambaBlock::new_seeded(cfg.clone(), dev, seed).expect("base");
+    let d = cfg.get_d_inner();
+    let mut map = HashMap::new();
+    for (name, t) in base.persistence_tensors() {
+        if name == "dt_bias" {
+            map.insert(name.to_string(), Tensor::from_vec(vec![0.0f32; d], &[d], dev).expect("dtb"));
+        } else {
+            map.insert(name.to_string(), t.clone());
+        }
+    }
+    MambaBlock::from_persistence_tensors(cfg.clone(), dev, map).expect("rebuild")
+}
+
 #[test]
 fn every_formerly_dead_weight_is_live() {
     let dev = Device::cpu();
-    let block = MambaBlock::new_seeded(cfg(), &dev, 42).expect("block");
+    let block = active_block(&cfg(), &dev, 42);
     let x = Tensor::randn_seeded(&[1, L, D_MODEL], &dev, 7).expect("x");
     let y0 = fwd(&block, &x);
 
@@ -216,7 +236,11 @@ fn training_loop_reduces_loss_and_moves_weights() {
     let target = teacher.forward(&x).expect("teacher").output.to_vec().expect("t");
 
     let mut block = scaled_block(&cfg, &dev, 5, 0.2);
-    let a_log_before = named_param(&block, "A_log");
+    // conv1d_weight is a seed-varying backbone weight (A_log/dt_bias/D
+    // are deterministic init, identical between teacher and student, so
+    // they carry little gradient here) — it must move if the SSM block
+    // (not just a head) is being trained.
+    let cw_before = named_param(&block, "conv1d_weight");
 
     let (lr, b1, b2, eps) = (0.03f32, 0.9f32, 0.999f32, 1e-8f32);
     let mut state: HashMap<String, (Vec<f32>, Vec<f32>)> = HashMap::new();
@@ -267,10 +291,10 @@ fn training_loop_reduces_loss_and_moves_weights() {
         last_loss < 0.5 * first_loss,
         "training did not reduce loss enough: {first_loss:.5} → {last_loss:.5}"
     );
-    let a_log_after = named_param(&block, "A_log");
+    let cw_after = named_param(&block, "conv1d_weight");
     assert!(
-        max_abs_diff(&a_log_before, &a_log_after) > 1e-3,
-        "backbone weight A_log did not move during training"
+        max_abs_diff(&cw_before, &cw_after) > 1e-3,
+        "backbone weight conv1d_weight did not move during training"
     );
 }