Mamba M3: trainable — Adam loop reduces loss and moves weights

Demonstrates end-to-end trainability of the real selective scan. A
self-contained Adam loop (the rtx-transformers AdamOptimizer has no
public gradient setter — the spec permits a bespoke loop) fits a teacher
block's output on a fixed input: forward → MSE → analytic backward →
Adam step → rebuild. Over 200 steps the loss drops >50% and the backbone
weight A_log moves, confirming gradients actually train the model (not
just the head). All 6 selective-scan tests green.

The production AdamOptimizer can be wired once it exposes a gradient
setter; the M2 backward already returns grads in its HashMap shape.

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>

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

@@ -189,6 +189,91 @@ fn analytic_gradients_match_finite_differences() {
     assert!(checked >= 30, "expected to check ≥30 grad elements, got {checked}");
 }
 
+fn named_param(block: &MambaBlock, name: &str) -> Vec<f32> {
+    block
+        .persistence_tensors()
+        .into_iter()
+        .find(|(nm, _)| *nm == name)
+        .map(|(_, t)| t.to_vec().expect("to_vec"))
+        .expect("param")
+}
+
+/// Milestone-3: a self-contained Adam loop drives the loss down and
+/// actually moves the backbone weights — i.e. the block is *trainable*.
+/// The target is a teacher block's output on the same input (a learnable
+/// regression at matched capacity), so a real gradient signal must shrink
+/// the loss substantially.
+#[test]
+fn training_loop_reduces_loss_and_moves_weights() {
+    let dev = Device::cpu();
+    let cfg = MambaConfig::new(4, 4, 3);
+    let ll = 4usize;
+    let dmodel = 4usize;
+
+    let teacher = scaled_block(&cfg, &dev, 99, 0.2);
+    let xvec: Vec<f32> = (0..ll * dmodel).map(|i| 0.5 * ((i as f32 * 0.41).sin())).collect();
+    let x = Tensor::from_vec(xvec, &[1, ll, dmodel], &dev).expect("x");
+    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");
+
+    let (lr, b1, b2, eps) = (0.03f32, 0.9f32, 0.999f32, 1e-8f32);
+    let mut state: HashMap<String, (Vec<f32>, Vec<f32>)> = HashMap::new();
+    let mut first_loss = 0.0f32;
+    let mut last_loss = 0.0f32;
+
+    for step in 1..=200i32 {
+        let out = block.forward(&x).expect("fwd").output.to_vec().expect("o");
+        let ne = out.len();
+        let mut dov = vec![0.0f32; ne];
+        let mut loss = 0.0f32;
+        for i in 0..ne {
+            let e = out[i] - target[i];
+            loss += e * e;
+            dov[i] = 2.0 * e / (ne as f32);
+        }
+        loss /= ne as f32;
+        if step == 1 {
+            first_loss = loss;
+        }
+        last_loss = loss;
+
+        let d_out = Tensor::from_vec(dov, &[1, ll, dmodel], &dev).expect("dov");
+        let grads = block.backward(&x, &d_out).expect("backward");
+
+        let mut map = HashMap::new();
+        for (name, t) in block.persistence_tensors() {
+            let p = t.to_vec().expect("p");
+            let g = grads.get(name).expect("g").to_vec().expect("gv");
+            let entry = state
+                .entry(name.to_string())
+                .or_insert_with(|| (vec![0.0f32; p.len()], vec![0.0f32; p.len()]));
+            let (m, v) = entry;
+            let mut np = vec![0.0f32; p.len()];
+            for i in 0..p.len() {
+                m[i] = b1 * m[i] + (1.0 - b1) * g[i];
+                v[i] = b2 * v[i] + (1.0 - b2) * g[i] * g[i];
+                let mhat = m[i] / (1.0 - b1.powi(step));
+                let vhat = v[i] / (1.0 - b2.powi(step));
+                np[i] = p[i] - lr * mhat / (vhat.sqrt() + eps);
+            }
+            map.insert(name.to_string(), Tensor::from_vec(np, t.shape().dims(), &dev).expect("np"));
+        }
+        block = MambaBlock::from_persistence_tensors(cfg.clone(), &dev, map).expect("rebuild");
+    }
+
+    assert!(
+        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");
+    assert!(
+        max_abs_diff(&a_log_before, &a_log_after) > 1e-3,
+        "backbone weight A_log did not move during training"
+    );
+}
+
 #[test]
 fn output_is_finite_and_non_constant() {
     let dev = Device::cpu();