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

@@ -0,0 +1,163 @@
+//! Phase-4a: the trained real selective-scan Mamba learns temporal
+//! structure that a memoryless model cannot.
+//!
+//! Task: next-token prediction on a multi-regime sequence. Each regime
+//! `r` has a fixed mean vector `μ_r`; the sequence visits regimes in
+//! blocks; `x[t] = μ_{r[t]} + noise`. To predict `x[t+1]` well *inside*
+//! a regime you must estimate `μ_{r[t]}` — which requires integrating
+//! recent history to average out the noise. A memoryless predictor
+//! (persistence: `x[t+1] ≈ x[t]`) is stuck at the full per-step noise;
+//! the SSM's state integrates, halving it.
+//!
+//! We assert the *trained* Mamba beats (a) the memoryless persistence
+//! baseline and (b) its own untrained self — i.e. training the real
+//! backbone (not a head) buys a genuine temporal-modeling win. This is
+//! the Phase-3/4 payoff: the replacement for the old "random Mamba 6.5%
+//! vs linear 48%" non-result.
+
+use std::collections::HashMap;
+
+use rtx_tensor::{Device, Tensor};
+use rtx_transformers::layers::mamba::{MambaBlock, MambaConfig};
+
+const DIM: usize = 16;
+const D_STATE: usize = 16;
+const D_CONV: usize = 4;
+const N_REGIMES: usize = 4;
+const BLOCK: usize = 8;
+const L: usize = N_REGIMES * BLOCK * 2; // two passes through the regime cycle
+const SIGMA: f32 = 0.3;
+
+/// Build a block with `dt_bias = 0` ⇒ Δ ≈ 0.69 so the scan is an active
+/// (leaky) integrator from the start, giving training a usable signal.
+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")
+}
+
+/// Deterministic multi-regime corpus, `[L, DIM]` row-major.
+fn corpus() -> Vec<f32> {
+    // SplitMix64 noise for reproducibility.
+    let mut s: u64 = 0xC0FFEE_1234_5678;
+    let mut noise = || -> f32 {
+        s = s.wrapping_add(0x9E37_79B9_7F4A_7C15);
+        let mut z = s;
+        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
+        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
+        let u = ((z ^ (z >> 31)) >> 11) as f32 / (1u64 << 53) as f32; // [0,1)
+        (u - 0.5) * 2.0 // ~U(-1,1)
+    };
+    let mu = |r: usize, d: usize| -> f32 {
+        ((r as f32) * 1.3 + (d as f32) * 0.5).sin() * 0.6
+    };
+    let mut x = vec![0.0f32; L * DIM];
+    for t in 0..L {
+        let r = (t / BLOCK) % N_REGIMES;
+        for d in 0..DIM {
+            x[t * DIM + d] = mu(r, d) + SIGMA * noise();
+        }
+    }
+    x
+}
+
+/// Mean-squared next-token prediction error of `block` on `x`, over
+/// `t ∈ [0, L-1)` (target `x[t+1]`). Returns (mse, d_out gradient).
+fn predict_mse(block: &MambaBlock, x: &Tensor, xv: &[f32], dev: &Device) -> (f32, Vec<f32>) {
+    let out = block.forward(x).expect("fwd").output.to_vec().expect("o");
+    let n = (L - 1) * DIM;
+    let mut mse = 0.0f32;
+    let mut dov = vec![0.0f32; L * DIM];
+    for t in 0..L - 1 {
+        for d in 0..DIM {
+            let pred = out[t * DIM + d];
+            let tgt = xv[(t + 1) * DIM + d];
+            let e = pred - tgt;
+            mse += e * e;
+            dov[t * DIM + d] = 2.0 * e / (n as f32);
+        }
+    }
+    let _ = dev;
+    (mse / n as f32, dov)
+}
+
+/// Memoryless baseline: predict `x[t+1] = x[t]`.
+fn persistence_mse(xv: &[f32]) -> f32 {
+    let n = (L - 1) * DIM;
+    let mut mse = 0.0f32;
+    for t in 0..L - 1 {
+        for d in 0..DIM {
+            let e = xv[t * DIM + d] - xv[(t + 1) * DIM + d];
+            mse += e * e;
+        }
+    }
+    mse / n as f32
+}
+
+#[test]
+fn trained_mamba_beats_memoryless_baseline_on_temporal_task() {
+    let dev = Device::cpu();
+    let cfg = MambaConfig::new(DIM, D_STATE, D_CONV);
+    let xv = corpus();
+    let x = Tensor::from_vec(xv.clone(), &[1, L, DIM], &dev).expect("x");
+
+    let baseline = persistence_mse(&xv);
+
+    let mut block = active_block(&cfg, &dev, 7);
+    let (untrained_mse, _) = predict_mse(&block, &x, &xv, &dev);
+
+    // Adam training loop (the rtx-transformers AdamOptimizer has no
+    // public gradient setter; a bespoke loop over the analytic-backward
+    // grad map is what the Phase-3 spec sanctions).
+    let (lr, b1, b2, eps) = (0.02f32, 0.9f32, 0.999f32, 1e-8f32);
+    let mut state: HashMap<String, (Vec<f32>, Vec<f32>)> = HashMap::new();
+    for step in 1..=400i32 {
+        let (_, dov) = predict_mse(&block, &x, &xv, &dev);
+        let d_out = Tensor::from_vec(dov, &[1, L, DIM], &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 (m, v) = state
+                .entry(name.to_string())
+                .or_insert_with(|| (vec![0.0f32; p.len()], vec![0.0f32; p.len()]));
+            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");
+    }
+
+    let (trained_mse, _) = predict_mse(&block, &x, &xv, &dev);
+    eprintln!(
+        "Phase-4a temporal learning: persistence(memoryless)={baseline:.4} \
+         untrained_mamba={untrained_mse:.4} trained_mamba={trained_mse:.4}"
+    );
+
+    // The trained SSM integrates history to de-noise the regime mean,
+    // beating the memoryless persistence baseline...
+    assert!(
+        trained_mse < baseline,
+        "trained Mamba ({trained_mse:.4}) did not beat memoryless persistence ({baseline:.4})"
+    );
+    // ...and training the real backbone genuinely improved it.
+    assert!(
+        trained_mse < untrained_mse,
+        "training did not improve the Mamba ({untrained_mse:.4} → {trained_mse:.4})"
+    );
+}