Merge pull request 'SMT E1+E3a/b: predictive-state teacher + cloned recurrent memory updater' (#5) from smt-set-encoder-teacher into main

crates/training/rtx-transformers/Cargo.toml

@@ -12,6 +12,10 @@ description = "Complete transformer training infrastructure with revolutionary q
 # Core RTX dependencies - enabled for autograd integration
 rtx-tensor = { workspace = true }
 rtx-autograd = { workspace = true }
+# Tape backend for the SMT predictive-state teacher (SetEncoderTeacher trains
+# on `Autodiff<CpuBackend>` — the gradient-checked real-backend path).
+rtx-backend = { workspace = true }
+rtx-backend-cpu = { path = "../../core/rtx-backend-cpu", version = "1.0.0" }
 
 # Essential dependencies
 anyhow = { workspace = true }

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

@@ -0,0 +1,351 @@
+//! `ClonedMemoryUpdater` — the recurrent memory updater behaviorally cloned
+//! from the SMT oracle (Supervised Memory Training, Phase E3b).
+//!
+//! Isola's *"Pre-training recurrent nets without recurrence"* trains the
+//! recurrent updater **without** backprop-through-time: the time-parallel
+//! teacher (the [`set_encoder_teacher`](super::set_encoder_teacher)) emits an
+//! oracle trajectory of predictive-state memories `M_1 … M_T`, and a recurrent
+//! cell `g(M_{t-1}, x_t) → M_t` is regressed onto that trajectory **one step at
+//! a time, teacher-forced on the oracle memory**. Because every training example
+//! feeds the *oracle's* previous memory (never the cell's own output) and the
+//! loss is a single-step MSE, there is no rollout to differentiate through — no
+//! BPTT.
+//!
+//! ## Architecture
+//! ```text
+//!   cat   = [M_{t-1} ‖ x_t]                  [d_mem + d_in]
+//!   h     = gelu(cat · W_in)                 [d_hidden]
+//!   M_t   = ρ·M_{t-1} + h · W_mem            [d_mem]   (leaky/contractive update)
+//!   x̂_{t+1} = M_t · W_read                   [d_in]    (predict-the-future head)
+//! ```
+//! The leaky update (`M_t = ρ·M_{t-1} + Δ`, `ρ < 1`) makes the cell learn the
+//! *increment* the new token induces while keeping the recurrence a contraction
+//! — so a free autoregressive rollout stays bounded and the (length-invariant)
+//! rule extrapolates past the teacher's training horizon.
+//!
+//! Trains on `Autodiff<CpuBackend>` with a self-contained deterministic Adam,
+//! mirroring [`set_encoder_teacher`](super::set_encoder_teacher). (The small
+//! tape-param boilerplate is intentionally duplicated rather than shared, to
+//! keep each SMT layer independently legible.)
+
+use rtx_autograd::autodiff::{Autodiff, AutodiffDevice, GradTensor, TensorId, backward_impl};
+use rtx_backend::{AutodiffBackend, Backend};
+use rtx_backend_cpu::CpuBackend;
+
+type Ad = Autodiff<CpuBackend>;
+
+/// Leaky memory decay `M_t = ρ·M_{t-1} + Δ`. With `ρ < 1` the recurrent state is
+/// a contraction, so a free autoregressive rollout stays bounded (`|M| ≤
+/// |Δ|_max / (1−ρ)`) instead of accumulating without limit — the difference
+/// between a usable rollout and a `1e9` blow-up.
+const DECAY: f32 = 0.9;
+
+/// Magnitude bound on the free-running memory state during autoregressive
+/// rollout (well outside the range a teacher-forced oracle memory ever reaches).
+const MEM_CLAMP: f32 = 4.0;
+
+/// Construction knobs for [`ClonedMemoryUpdater`].
+#[derive(Clone, Copy, Debug)]
+pub struct ClonedUpdaterConfig {
+    /// Token width (the latent dimension of `x` and the readout).
+    pub d_in: usize,
+    /// Memory size `|M|` — must match the oracle teacher's `d_mem`.
+    pub d_mem: usize,
+    /// Hidden width of the update MLP.
+    pub d_hidden: usize,
+    /// Deterministic init seed.
+    pub seed: u64,
+}
+
+impl ClonedUpdaterConfig {
+    /// Standard config.
+    #[must_use]
+    pub const fn new(d_in: usize, d_mem: usize, d_hidden: usize, seed: u64) -> Self {
+        Self {
+            d_in,
+            d_mem,
+            d_hidden,
+            seed,
+        }
+    }
+}
+
+/// Deterministic init in `[-scale, scale)` via a SplitMix64-style LCG.
+fn lcg_init(n: usize, scale: f32, seed: u64) -> Vec<f32> {
+    let mut s = seed.wrapping_add(0x9E37_79B9_7F4A_7C15);
+    (0..n)
+        .map(|_| {
+            s = s
+                .wrapping_mul(6_364_136_223_846_793_005)
+                .wrapping_add(1_442_695_040_888_963_407);
+            let u = ((s >> 33) as f32) / ((1u64 << 31) as f32);
+            (u * 2.0 - 1.0) * scale
+        })
+        .collect()
+}
+
+/// A trainable parameter matrix with its own Adam moments (host-side).
+struct Param {
+    data: Vec<f32>,
+    rows: usize,
+    cols: usize,
+    m: Vec<f32>,
+    v: Vec<f32>,
+}
+
+impl Param {
+    fn new(rows: usize, cols: usize, scale: f32, seed: u64) -> Self {
+        let n = rows * cols;
+        Self {
+            data: lcg_init(n, scale, seed),
+            rows,
+            cols,
+            m: vec![0.0; n],
+            v: vec![0.0; n],
+        }
+    }
+
+    fn shape(&self) -> [usize; 2] {
+        [self.rows, self.cols]
+    }
+
+    fn leaf(
+        &self,
+        dev: &AutodiffDevice<CpuBackend>,
+    ) -> (<Ad as Backend>::TensorPrimitive<2>, usize) {
+        let t = Ad::require_grad(Ad::from_data(&self.data, self.shape(), dev));
+        let id = t.id().0;
+        (t, id)
+    }
+
+    fn adam(&mut self, grad: &[f32], lr: f32, t: u32) {
+        const B1: f32 = 0.9;
+        const B2: f32 = 0.999;
+        const EPS: f32 = 1e-8;
+        let bc1 = 1.0 - B1.powi(t as i32);
+        let bc2 = 1.0 - B2.powi(t as i32);
+        for ((w, g), (m, v)) in self
+            .data
+            .iter_mut()
+            .zip(grad.iter())
+            .zip(self.m.iter_mut().zip(self.v.iter_mut()))
+        {
+            *m = B1 * *m + (1.0 - B1) * g;
+            *v = B2 * *v + (1.0 - B2) * g * g;
+            *w -= lr * (*m / bc1) / ((*v / bc2).sqrt() + EPS);
+        }
+    }
+}
+
+fn grad_of(storage: &rtx_autograd::autodiff::GradientStorage<CpuBackend>, id: usize) -> Vec<f32> {
+    match storage
+        .get(TensorId(id))
+        .unwrap_or_else(|| panic!("no gradient for leaf {id}"))
+    {
+        GradTensor::D2(t) => t.to_vec(),
+        other => panic!("expected D2 gradient, got {other:?}"),
+    }
+}
+
+/// The cloned recurrent memory updater + predict-the-future readout.
+pub struct ClonedMemoryUpdater {
+    cfg: ClonedUpdaterConfig,
+    w_in: Param,   // [d_mem + d_in, d_hidden]
+    w_mem: Param,  // [d_hidden, d_mem]
+    w_read: Param, // [d_mem, d_in]
+    step: u32,
+    dev: AutodiffDevice<CpuBackend>,
+}
+
+impl ClonedMemoryUpdater {
+    /// Build with deterministic, seeded weights.
+    #[must_use]
+    pub fn new(cfg: ClonedUpdaterConfig) -> Self {
+        let s = |i: u64| cfg.seed.wrapping_mul(0x100_0000_01b3).wrapping_add(i);
+        let sc = |fan: usize| 1.0 / (fan as f32).sqrt();
+        Self {
+            w_in: Param::new(
+                cfg.d_mem + cfg.d_in,
+                cfg.d_hidden,
+                sc(cfg.d_mem + cfg.d_in),
+                s(1),
+            ),
+            // Small last-layer init so the residual update starts near identity.
+            w_mem: Param::new(cfg.d_hidden, cfg.d_mem, 0.1 * sc(cfg.d_hidden), s(2)),
+            w_read: Param::new(cfg.d_mem, cfg.d_in, sc(cfg.d_mem), s(3)),
+            step: 0,
+            dev: AutodiffDevice::<CpuBackend>::default(),
+            cfg,
+        }
+    }
+
+    /// Off-tape recurrent step: `g(prev_mem, x) → new_mem` (`[d_mem]`).
+    ///
+    /// # Panics
+    /// Panics if `prev_mem.len() != d_mem` or `x.len() != d_in`.
+    #[must_use]
+    pub fn step(&self, prev_mem: &[f32], x: &[f32]) -> Vec<f32> {
+        assert_eq!(prev_mem.len(), self.cfg.d_mem);
+        assert_eq!(x.len(), self.cfg.d_in);
+        let dh = self.cfg.d_hidden;
+        let dm = self.cfg.d_mem;
+        // h = gelu([prev_mem ‖ x] · W_in)
+        let mut h = vec![0.0f32; dh];
+        for (k, hk) in h.iter_mut().enumerate() {
+            let mut acc = 0.0f32;
+            for (i, &pm) in prev_mem.iter().enumerate() {
+                acc += pm * self.w_in.data[i * dh + k];
+            }
+            for (i, &xi) in x.iter().enumerate() {
+                acc += xi * self.w_in.data[(dm + i) * dh + k];
+            }
+            *hk = gelu_scalar(acc);
+        }
+        // new_mem = ρ·prev_mem + h · W_mem  (leaky/contractive update)
+        let mut out: Vec<f32> = prev_mem.iter().map(|&p| DECAY * p).collect();
+        for (k, &hk) in h.iter().enumerate() {
+            for (j, oj) in out.iter_mut().enumerate().take(dm) {
+                *oj += hk * self.w_mem.data[k * dm + j];
+            }
+        }
+        // Rollout safety clamp: one-step BC has no signal against autoregressive
+        // drift (that is E4/DAgger's job), so bound the free-running state to
+        // keep a long rollout finite. Training is teacher-forced on the bounded
+        // oracle memory, so this never binds during training.
+        for o in &mut out {
+            *o = o.clamp(-MEM_CLAMP, MEM_CLAMP);
+        }
+        out
+    }
+
+    /// Off-tape readout: predict the next latent from a memory (`[d_in]`).
+    ///
+    /// # Panics
+    /// Panics if `mem.len() != d_mem`.
+    #[must_use]
+    pub fn predict(&self, mem: &[f32]) -> Vec<f32> {
+        assert_eq!(mem.len(), self.cfg.d_mem);
+        let din = self.cfg.d_in;
+        let mut out = vec![0.0f32; din];
+        for (j, oj) in out.iter_mut().enumerate() {
+            let mut acc = 0.0f32;
+            for (i, &mi) in mem.iter().enumerate() {
+                acc += mi * self.w_read.data[i * din + j];
+            }
+            *oj = acc;
+        }
+        out
+    }
+
+    /// One behavioral-cloning step (no BPTT, teacher-forced on the oracle
+    /// memory). Trains the updater to map `(prev_mem, x) → target_mem` and the
+    /// readout to map the produced memory → `target_next`. Returns the total
+    /// loss before the update.
+    ///
+    /// # Panics
+    /// Panics on any shape mismatch.
+    pub fn train_step(
+        &mut self,
+        prev_mem: &[f32],
+        x: &[f32],
+        target_mem: &[f32],
+        target_next: &[f32],
+    ) -> f32 {
+        let dm = self.cfg.d_mem;
+        let din = self.cfg.d_in;
+        assert_eq!(prev_mem.len(), dm);
+        assert_eq!(x.len(), din);
+        assert_eq!(target_mem.len(), dm);
+        assert_eq!(target_next.len(), din);
+
+        let (w_in, in_id) = self.w_in.leaf(&self.dev);
+        let (w_mem, mem_id) = self.w_mem.leaf(&self.dev);
+        let (w_read, read_id) = self.w_read.leaf(&self.dev);
+
+        // cat = [prev_mem ‖ x]  → [1, d_mem + d_in]
+        let mut cat = Vec::with_capacity(dm + din);
+        cat.extend_from_slice(prev_mem);
+        cat.extend_from_slice(x);
+        let cat_t = Ad::from_data(&cat, [1, dm + din], &self.dev);
+        let prev_t = Ad::from_data(prev_mem, [1, dm], &self.dev);
+
+        let h = Ad::gelu(Ad::matmul(cat_t, w_in)); // [1, d_hidden]
+        let delta = Ad::matmul(h, w_mem); // [1, d_mem]
+        let decay_t = Ad::from_data(&vec![DECAY; dm], [1, dm], &self.dev);
+        let new_mem = Ad::add(Ad::mul(prev_t, decay_t), delta); // leaky update
+        let pred = Ad::matmul(new_mem.clone(), w_read); // [1, d_in]
+
+        let tgt_mem_t = Ad::from_data(target_mem, [1, dm], &self.dev);
+        let tgt_next_t = Ad::from_data(target_next, [1, din], &self.dev);
+        let mem_err = Ad::sub(new_mem, tgt_mem_t);
+        let read_err = Ad::sub(pred, tgt_next_t);
+        let loss = Ad::add(
+            Ad::sum(Ad::mul(mem_err.clone(), mem_err)),
+            Ad::sum(Ad::mul(read_err.clone(), read_err)),
+        );
+        let loss_val = Ad::to_data(&loss)[0];
+
+        let storage = backward_impl(
+            &loss,
+            Some(GradTensor::from_d1(CpuBackend::ones(
+                [1],
+                &<CpuBackend as Backend>::Device::default(),
+            ))),
+        )
+        .expect("backward");
+
+        self.step += 1;
+        let t = self.step;
+        const LR: f32 = 0.01;
+        self.w_in.adam(&grad_of(&storage, in_id), LR, t);
+        self.w_mem.adam(&grad_of(&storage, mem_id), LR, t);
+        self.w_read.adam(&grad_of(&storage, read_id), LR, t);
+        loss_val
+    }
+}
+
+/// Exact-ish GELU matching the tape's `Ad::gelu` (tanh approximation), so the
+/// off-tape `step`/`predict` agree with the trained graph.
+fn gelu_scalar(x: f32) -> f32 {
+    const C: f32 = 0.797_884_56; // sqrt(2/pi)
+    0.5 * x * (1.0 + (C * (x + 0.044_715 * x * x * x)).tanh())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn updater_learns_a_one_step_transition() {
+        // Teach the cell a fixed (prev_mem, x) → target_mem map + readout.
+        let cfg = ClonedUpdaterConfig::new(4, 6, 16, 1);
+        let mut g = ClonedMemoryUpdater::new(cfg);
+        let prev = lcg_init(6, 1.0, 10);
+        let x = lcg_init(4, 1.0, 11);
+        let target_mem = lcg_init(6, 0.5, 12);
+        let target_next = lcg_init(4, 0.5, 13);
+
+        let first = g.train_step(&prev, &x, &target_mem, &target_next);
+        let mut last = first;
+        for _ in 0..400 {
+            last = g.train_step(&prev, &x, &target_mem, &target_next);
+            assert!(last.is_finite());
+        }
+        assert!(
+            last < first * 0.1,
+            "updater did not learn: {first:.4} -> {last:.4}"
+        );
+    }
+
+    #[test]
+    fn off_tape_step_matches_trained_graph_shape() {
+        let cfg = ClonedUpdaterConfig::new(4, 6, 8, 2);
+        let g = ClonedMemoryUpdater::new(cfg);
+        let prev = vec![0.0f32; 6];
+        let x = lcg_init(4, 1.0, 5);
+        let new_mem = g.step(&prev, &x);
+        assert_eq!(new_mem.len(), 6);
+        assert_eq!(g.predict(&new_mem).len(), 4);
+        assert!(new_mem.iter().all(|v| v.is_finite()));
+    }
+}

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

@@ -0,0 +1,274 @@
+//! `MambaRecurrence` — the single-timestep recurrent view of a [`MambaBlock`]
+//! (Supervised Memory Training, Phase E3).
+//!
+//! The full-sequence [`MambaBlock::forward`] runs the selective scan inline over
+//! a whole window. To behaviorally-clone Mamba as a *memory updater*
+//! `f(state, x_t) → (state', y_t)` (the SMT core — no backprop-through-time), we
+//! need the recurrence exposed one step at a time. This module reconstructs
+//! exactly that recurrence from a block's persisted weights, with **bit-level
+//! fidelity** to the full forward: stepping a sequence through `step` and
+//! concatenating the per-step outputs reproduces `forward` (pinned by
+//! `tests/mamba_step_equivalence.rs`).
+//!
+//! The recurrent state is `(h, conv)`:
+//! - `h ∈ ℝ^{d_inner · d_state}` — the selective-scan SSM state `h[j·n + k]`.
+//! - `conv ∈ ℝ^{(d_conv − 1) · d_inner}` — the causal depthwise-conv ring of the
+//!   last `d_conv − 1` pre-activation inputs (zero-initialized = the forward's
+//!   left zero-padding).
+//!
+//! Dimensions are derived from the persisted tensor shapes, so no access to the
+//! block's private config is needed and `mamba.rs` is left untouched.
+
+use crate::Result;
+use crate::layers::mamba::MambaBlock;
+
+/// Numerically-stable SiLU `x·σ(x)` — identical to the one in `mamba.rs` so the
+/// stepwise path matches the forward bit-for-bit.
+#[inline]
+fn silu(x: f32) -> f32 {
+    let s = if x >= 0.0 {
+        1.0 / (1.0 + (-x).exp())
+    } else {
+        let e = x.exp();
+        e / (1.0 + e)
+    };
+    x * s
+}
+
+/// Numerically-stable softplus `ln(1 + eˣ)` — identical to the one in `mamba.rs`.
+#[inline]
+fn softplus(x: f32) -> f32 {
+    x.max(0.0) + (1.0 + (-x.abs()).exp()).ln()
+}
+
+/// The recurrent state carried between [`MambaRecurrence::step`] calls.
+#[derive(Clone, Debug)]
+pub struct MambaState {
+    /// Selective-scan SSM state, row-major `[d_inner, d_state]` as `h[j*n + k]`.
+    h: Vec<f32>,
+    /// Causal-conv ring of the last `d_conv − 1` pre-activation inputs,
+    /// row-major `[d_conv - 1, d_inner]` (index 0 = oldest).
+    conv: Vec<f32>,
+}
+
+impl MambaState {
+    /// Read-only view of the SSM state `h` (`[d_inner · d_state]`). This is the
+    /// "memory" the SMT cloning supervises.
+    #[must_use]
+    pub fn hidden(&self) -> &[f32] {
+        &self.h
+    }
+}
+
+/// A `MambaBlock` reconstituted as an explicit step-by-step recurrence over host
+/// `f32` weights.
+pub struct MambaRecurrence {
+    d_model: usize,
+    d_inner: usize,
+    d_state: usize,
+    dt_rank: usize,
+    d_conv: usize,
+    in_proj: Vec<f32>,  // [d_model, 2*d_inner]
+    conv_w: Vec<f32>,   // [d_inner, d_conv]
+    conv_b: Vec<f32>,   // [d_inner] (zeros if the block had no conv bias)
+    a_log: Vec<f32>,    // [d_inner, d_state]
+    x_proj: Vec<f32>,   // [d_inner, dt_rank + 2*d_state]
+    dt_proj: Vec<f32>,  // [dt_rank, d_inner]
+    dt_bias: Vec<f32>,  // [d_inner]
+    d_skip: Vec<f32>,   // [d_inner]
+    out_proj: Vec<f32>, // [d_inner, d_model]
+}
+
+impl MambaRecurrence {
+    /// Reconstruct the recurrence from a block's persisted weights.
+    ///
+    /// # Errors
+    /// Propagates any tensor `to_vec` failure.
+    pub fn from_block(block: &MambaBlock) -> Result<Self> {
+        let mut in_proj = Vec::new();
+        let mut conv_w = Vec::new();
+        let mut conv_b = None;
+        let mut a_log = Vec::new();
+        let mut x_proj = Vec::new();
+        let mut dt_proj = Vec::new();
+        let mut dt_bias = Vec::new();
+        let mut d_skip = Vec::new();
+        let mut out_proj = Vec::new();
+        // Shapes we recover dimensions from.
+        let mut in_proj_shape = Vec::new();
+        let mut conv_shape = Vec::new();
+        let mut a_shape = Vec::new();
+        let mut x_proj_cols = 0usize;
+
+        for (name, tensor) in block.persistence_tensors() {
+            let dims = tensor.shape().dims().to_vec();
+            let data = tensor.to_vec()?;
+            match name {
+                "in_proj" => {
+                    in_proj_shape = dims;
+                    in_proj = data;
+                }
+                "conv1d_weight" => {
+                    conv_shape = dims;
+                    conv_w = data;
+                }
+                "conv1d_bias" => conv_b = Some(data),
+                "A_log" => {
+                    a_shape = dims;
+                    a_log = data;
+                }
+                "x_proj" => {
+                    x_proj_cols = dims[1];
+                    x_proj = data;
+                }
+                "dt_proj" => dt_proj = data,
+                "dt_bias" => dt_bias = data,
+                "D" => d_skip = data,
+                "out_proj" => out_proj = data,
+                _ => {}
+            }
+        }
+
+        let d_model = in_proj_shape[0];
+        let d_inner = in_proj_shape[1] / 2;
+        let d_conv = conv_shape[2]; // [d_inner, 1, d_conv]
+        let d_state = a_shape[1];
+        let dt_rank = x_proj_cols - 2 * d_state;
+        let conv_b = conv_b.unwrap_or_else(|| vec![0.0f32; d_inner]);
+
+        Ok(Self {
+            d_model,
+            d_inner,
+            d_state,
+            dt_rank,
+            d_conv,
+            in_proj,
+            conv_w,
+            conv_b,
+            a_log,
+            x_proj,
+            dt_proj,
+            dt_bias,
+            d_skip,
+            out_proj,
+        })
+    }
+
+    /// Input/output token width (`d_model`).
+    #[must_use]
+    pub const fn d_model(&self) -> usize {
+        self.d_model
+    }
+
+    /// Size of the recurrent SSM state `h` (`d_inner · d_state`).
+    #[must_use]
+    pub const fn state_size(&self) -> usize {
+        self.d_inner * self.d_state
+    }
+
+    /// A fresh zero state (matches the forward's zero conv-padding + zero `h`).
+    #[must_use]
+    pub fn init_state(&self) -> MambaState {
+        MambaState {
+            h: vec![0.0f32; self.d_inner * self.d_state],
+            conv: vec![0.0f32; self.d_conv.saturating_sub(1) * self.d_inner],
+        }
+    }
+
+    /// Advance one timestep: consume `x_t` (`[d_model]`), mutate `state`, and
+    /// return the output `y_t` (`[d_model]`). Bit-identical to the corresponding
+    /// timestep of [`MambaBlock::forward`].
+    ///
+    /// # Panics
+    /// Panics if `x_t.len() != d_model`.
+    #[must_use]
+    pub fn step(&self, x_t: &[f32], state: &mut MambaState) -> Vec<f32> {
+        let d = self.d_inner;
+        let n = self.d_state;
+        let kc = self.d_conv;
+        let dbc = self.dt_rank + 2 * n;
+        assert_eq!(x_t.len(), self.d_model, "x_t must be [d_model]");
+
+        // (1) in_proj → x_in (scan branch) + z (gate branch).
+        let mut x_in = vec![0.0f32; d];
+        let mut z = vec![0.0f32; d];
+        for (j, (xij, zj)) in x_in.iter_mut().zip(z.iter_mut()).enumerate() {
+            let mut sx = 0.0f32;
+            let mut sz = 0.0f32;
+            for (m, &xm) in x_t.iter().enumerate() {
+                sx += xm * self.in_proj[m * (2 * d) + j];
+                sz += xm * self.in_proj[m * (2 * d) + d + j];
+            }
+            *xij = sx;
+            *zj = sz;
+        }
+
+        // (2-3) causal depthwise conv over the last kc pre-activation inputs,
+        // then SiLU. window[kk] = conv-ring for kk < kc-1, else the current x_in.
+        let mut u = vec![0.0f32; d];
+        for j in 0..d {
+            let mut acc = self.conv_b[j];
+            for kk in 0..kc {
+                let v = if kk + 1 < kc {
+                    state.conv[kk * d + j]
+                } else {
+                    x_in[j]
+                };
+                acc += self.conv_w[j * kc + kk] * v;
+            }
+            u[j] = silu(acc);
+        }
+        // Slide the conv ring: drop the oldest, append the current x_in.
+        if kc > 1 {
+            for k in 0..(kc - 2) {
+                for j in 0..d {
+                    state.conv[k * d + j] = state.conv[(k + 1) * d + j];
+                }
+            }
+            for j in 0..d {
+                state.conv[(kc - 2) * d + j] = x_in[j];
+            }
+        }
+
+        // (4) x_proj → (dt_logits, B, C).
+        let mut xdbl = vec![0.0f32; dbc];
+        for (q, xq) in xdbl.iter_mut().enumerate() {
+            let mut s = 0.0f32;
+            for (j, &uj) in u.iter().enumerate() {
+                s += uj * self.x_proj[j * dbc + q];
+            }
+            *xq = s;
+        }
+
+        // (5) delta = softplus(dt_bias + dt_logits @ dt_proj).
+        let mut delta = vec![0.0f32; d];
+        for (j, dj) in delta.iter_mut().enumerate() {
+            let mut s = self.dt_bias[j];
+            for r in 0..self.dt_rank {
+                s += xdbl[r] * self.dt_proj[r * d + j];
+            }
+            *dj = softplus(s);
+        }
+
+        // (6-9) selective scan + skip + gate + out_proj.
+        let mut out = vec![0.0f32; self.d_model];
+        for j in 0..d {
+            let dj = delta[j];
+            let uj = u[j];
+            let mut y = self.d_skip[j] * uj;
+            for k in 0..n {
+                let a = -self.a_log[j * n + k].exp();
+                let da = (dj * a).exp();
+                let dbu = dj * xdbl[self.dt_rank + k] * uj; // dj * B[k] * uj
+                let hv = da * state.h[j * n + k] + dbu;
+                state.h[j * n + k] = hv;
+                y += xdbl[self.dt_rank + n + k] * hv; // C[k] * h'
+            }
+            let yg = y * silu(z[j]);
+            for m in 0..self.d_model {
+                out[m] += yg * self.out_proj[j * self.d_model + m];
+            }
+        }
+        out
+    }
+}

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

@@ -55,8 +55,17 @@ pub mod mixture_of_experts;
 
 // State-space models (Mamba)
 pub mod mamba;
+pub mod mamba_step;
 pub mod metal_mamba;
 
+// SMT predictive-state teacher (the time-parallel oracle for Supervised
+// Memory Training — see `set_encoder_teacher` for the method).
+pub mod set_encoder_teacher;
+
+// SMT cloned recurrent memory updater (behaviorally cloned from the oracle
+// trajectory, no BPTT — see `cloned_memory_updater`).
+pub mod cloned_memory_updater;
+
 // Ring attention for long context
 pub mod ring_attention;
 // pub mod mamba_integration;

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

@@ -0,0 +1,478 @@
+//! `SetEncoderTeacher` — the SMT predictive-state oracle (Supervised Memory
+//! Training, Phase E1).
+//!
+//! Phillip Isola's *"Pre-training recurrent nets without recurrence"* argues
+//! that a recurrent memory updater should be trained **without** backprop-
+//! through-time by *behaviorally cloning* the trajectory of an oracle's
+//! **predictive-state** memories — the optimal fixed-size sufficient statistic
+//! of the past for predicting the future. This struct is that oracle: a
+//! permutation-invariant **set encoder** that maps a window of past tokens to a
+//! fixed-size memory `M`, trained with a predict-the-future objective so that
+//! `M` becomes exactly such a sufficient statistic.
+//!
+//! ## Architecture (the shape proven in `tape_train_smoke.rs`)
+//! ```text
+//!   e      = X · W_embed                      [L, d_model]   (token embedding)
+//!   q,k,v  = e·W_q,  e·W_k,  e·W_v            [L, d_model]   (self-attention)
+//!   ctx    = softmax(q·kᵀ) · v                [L, d_model]
+//!   h      = e + ctx                          [L, d_model]   (residual)
+//!   pooled = meanₗ(h)                         [1, d_model]   (set pooling)
+//!   M      = pooled · W_mem                   [1, d_mem]     (fixed-size memory)
+//!   ŷ      = M · W_dec                        [1, d_in]      (predict-the-future)
+//! ```
+//! The embedding `e` fans out to Q/K/V/residual (four uses), exercising the
+//! autograd tape's fan-out gradient accumulation. Every op used (matmul,
+//! transpose, softmax, add, mul, sub, sum) is finite-difference gradient-checked
+//! in `tape_cpu_gradcheck.rs`; the whole network is trained end-to-end on
+//! `Autodiff<CpuBackend>` — the real, gradient-correct backend (not the
+//! shape-only `MockBackend`).
+//!
+//! ## Why no BPTT here
+//! The teacher is **time-parallel**: one window → one memory, no recurrence to
+//! unroll. Phase E2 runs it over every prefix to emit an oracle trajectory; a
+//! recurrent Mamba updater is then cloned against that trajectory one step at a
+//! time (Phase E3), so the *recurrent* net never needs BPTT.
+//!
+//! Persistence and provenance live with the caller (omni-think's
+//! `PredictiveStateTeacher` facade), which already owns the safetensors +
+//! BLAKE3 sealing path; this struct exposes raw parameter access via
+//! [`SetEncoderTeacher::named_params`] / [`SetEncoderTeacher::set_named_params`].
+
+use rtx_autograd::autodiff::{Autodiff, AutodiffDevice, GradTensor, TensorId, backward_impl};
+use rtx_backend::{AutodiffBackend, Backend};
+use rtx_backend_cpu::CpuBackend;
+
+/// Tape backend the teacher trains on: reverse-mode autodiff over the CPU
+/// backend (the gradient-checked real-backend path).
+type Ad = Autodiff<CpuBackend>;
+
+/// Construction knobs for [`SetEncoderTeacher`].
+#[derive(Clone, Copy, Debug)]
+pub struct SetEncoderConfig {
+    /// Token width — the latent dimension of each input/predicted token.
+    pub d_in: usize,
+    /// Internal model width of the attention encoder.
+    pub d_model: usize,
+    /// Fixed memory size `|M|`. This is the compressed sufficient statistic the
+    /// downstream recurrent updater will be cloned to reproduce.
+    pub d_mem: usize,
+    /// Deterministic init seed. Two teachers built with the same `(config,
+    /// seed)` are bit-identical, so training is reproducible.
+    pub seed: u64,
+}
+
+impl SetEncoderConfig {
+    /// Standard config for a `d_in`-wide latent stream.
+    #[must_use]
+    pub const fn new(d_in: usize, d_model: usize, d_mem: usize, seed: u64) -> Self {
+        Self {
+            d_in,
+            d_model,
+            d_mem,
+            seed,
+        }
+    }
+}
+
+/// One trainable parameter matrix held on the host as a flat row-major `Vec`,
+/// carrying its own Adam moment estimates. A fresh grad-tracked tape leaf is
+/// minted every step (define-by-run); the gradient read back from the tape
+/// drives an Adam update applied here on the host.
+struct Param {
+    data: Vec<f32>,
+    rows: usize,
+    cols: usize,
+    /// First-moment (mean) estimate, same length as `data`.
+    m: Vec<f32>,
+    /// Second-moment (uncentred variance) estimate, same length as `data`.
+    v: Vec<f32>,
+}
+
+impl Param {
+    fn new(rows: usize, cols: usize, scale: f32, seed: u64) -> Self {
+        let n = rows * cols;
+        Self {
+            data: lcg_init(n, scale, seed),
+            rows,
+            cols,
+            m: vec![0.0; n],
+            v: vec![0.0; n],
+        }
+    }
+
+    fn shape(&self) -> [usize; 2] {
+        [self.rows, self.cols]
+    }
+
+    /// A fresh grad-tracked leaf for this step plus its tape id.
+    fn leaf(
+        &self,
+        dev: &AutodiffDevice<CpuBackend>,
+    ) -> (<Ad as Backend>::TensorPrimitive<2>, usize) {
+        let t = Ad::require_grad(Ad::from_data(&self.data, self.shape(), dev));
+        let id = t.id().0;
+        (t, id)
+    }
+
+    /// Adam update from `grad` (same layout as `data`). `t` is the 1-based
+    /// global step used for bias correction.
+    fn adam(&mut self, grad: &[f32], lr: f32, t: u32) {
+        const B1: f32 = 0.9;
+        const B2: f32 = 0.999;
+        const EPS: f32 = 1e-8;
+        let bc1 = 1.0 - B1.powi(t as i32);
+        let bc2 = 1.0 - B2.powi(t as i32);
+        for ((w, g), (m, v)) in self
+            .data
+            .iter_mut()
+            .zip(grad.iter())
+            .zip(self.m.iter_mut().zip(self.v.iter_mut()))
+        {
+            *m = B1 * *m + (1.0 - B1) * g;
+            *v = B2 * *v + (1.0 - B2) * g * g;
+            let m_hat = *m / bc1;
+            let v_hat = *v / bc2;
+            *w -= lr * m_hat / (v_hat.sqrt() + EPS);
+        }
+    }
+}
+
+/// Deterministic init in `[-scale, scale)` via a SplitMix64-style LCG — avoids
+/// an all-zeros start (which gives zero gradients through a linear layer).
+fn lcg_init(n: usize, scale: f32, seed: u64) -> Vec<f32> {
+    let mut s = seed.wrapping_add(0x9E37_79B9_7F4A_7C15);
+    (0..n)
+        .map(|_| {
+            s = s
+                .wrapping_mul(6_364_136_223_846_793_005)
+                .wrapping_add(1_442_695_040_888_963_407);
+            let u = ((s >> 33) as f32) / ((1u64 << 31) as f32); // [0, 1)
+            (u * 2.0 - 1.0) * scale
+        })
+        .collect()
+}
+
+/// The predictive-state teacher: a trainable attention set-encoder producing a
+/// fixed-size memory `M`, with a decoder head supervised by predict-the-future.
+pub struct SetEncoderTeacher {
+    cfg: SetEncoderConfig,
+    w_embed: Param, // [d_in, d_model]
+    wq: Param,      // [d_model, d_model]
+    wk: Param,      // [d_model, d_model]
+    wv: Param,      // [d_model, d_model]
+    w_mem: Param,   // [d_model, d_mem]
+    w_dec: Param,   // [d_mem, d_in]
+    /// Global Adam step counter (1-based after the first update).
+    step: u32,
+    dev: AutodiffDevice<CpuBackend>,
+}
+
+/// Forward intermediates needed to recover `M` (and predictions) off-tape.
+struct Forward {
+    /// Fixed-size memory `M` as a `[d_mem]` host vector.
+    memory: Vec<f32>,
+    /// Predicted next token `ŷ` as a `[d_in]` host vector.
+    prediction: Vec<f32>,
+}
+
+impl SetEncoderTeacher {
+    /// Build a teacher with deterministic, seeded weights.
+    #[must_use]
+    pub fn new(cfg: SetEncoderConfig) -> Self {
+        // Per-tensor seeds derived from the base seed so no two matrices share
+        // an init stream.
+        let s = |i: u64| cfg.seed.wrapping_mul(0x100_0000_01b3).wrapping_add(i);
+        // 1/sqrt(fan_in)-style scales keep activations from blowing up.
+        let sc = |fan_in: usize| 1.0 / (fan_in as f32).sqrt();
+        Self {
+            w_embed: Param::new(cfg.d_in, cfg.d_model, sc(cfg.d_in), s(1)),
+            wq: Param::new(cfg.d_model, cfg.d_model, sc(cfg.d_model), s(2)),
+            wk: Param::new(cfg.d_model, cfg.d_model, sc(cfg.d_model), s(3)),
+            wv: Param::new(cfg.d_model, cfg.d_model, sc(cfg.d_model), s(4)),
+            w_mem: Param::new(cfg.d_model, cfg.d_mem, sc(cfg.d_model), s(5)),
+            w_dec: Param::new(cfg.d_mem, cfg.d_in, sc(cfg.d_mem), s(6)),
+            step: 0,
+            dev: AutodiffDevice::<CpuBackend>::default(),
+            cfg,
+        }
+    }
+
+    /// The teacher's configuration.
+    #[must_use]
+    pub const fn config(&self) -> SetEncoderConfig {
+        self.cfg
+    }
+
+    /// Encode a window of `l` tokens (`window` is row-major `[l, d_in]`) into the
+    /// fixed-size predictive-state memory `M` (`[d_mem]`).
+    ///
+    /// # Panics
+    /// Panics if `window.len() != l * d_in`.
+    #[must_use]
+    pub fn encode(&self, window: &[f32], l: usize) -> Vec<f32> {
+        self.run(window, l).memory
+    }
+
+    /// Predict the next token from a window of `l` tokens. Returns `[d_in]`.
+    ///
+    /// # Panics
+    /// Panics if `window.len() != l * d_in`.
+    #[must_use]
+    pub fn predict(&self, window: &[f32], l: usize) -> Vec<f32> {
+        self.run(window, l).prediction
+    }
+
+    /// One training step on a single `(window, target)` pair. `window` is
+    /// row-major `[l, d_in]`; `target` is the `[d_in]` next token. Returns the
+    /// sum-of-squared-error loss *before* the update.
+    ///
+    /// # Panics
+    /// Panics if `window.len() != l * d_in` or `target.len() != d_in`.
+    pub fn train_step(&mut self, window: &[f32], l: usize, target: &[f32]) -> f32 {
+        let din = self.cfg.d_in;
+        assert_eq!(window.len(), l * din, "window must be [l, d_in]");
+        assert_eq!(target.len(), din, "target must be [d_in]");
+
+        let (we, we_id) = self.w_embed.leaf(&self.dev);
+        let (q_w, q_id) = self.wq.leaf(&self.dev);
+        let (k_w, k_id) = self.wk.leaf(&self.dev);
+        let (v_w, v_id) = self.wv.leaf(&self.dev);
+        let (mem_w, mem_id) = self.w_mem.leaf(&self.dev);
+        let (dec_w, dec_id) = self.w_dec.leaf(&self.dev);
+
+        let loss = self.forward_graph(window, l, target, &we, &q_w, &k_w, &v_w, &mem_w, &dec_w);
+        let loss_val = Ad::to_data(&loss)[0];
+
+        let storage = backward_impl(
+            &loss,
+            Some(GradTensor::from_d1(CpuBackend::ones(
+                [1],
+                &<CpuBackend as Backend>::Device::default(),
+            ))),
+        )
+        .expect("backward");
+
+        self.step += 1;
+        let t = self.step;
+        // Adam learning rate — modest and fixed; the host loop owns scheduling.
+        const LR: f32 = 0.01;
+        self.w_embed.adam(&grad_of(&storage, we_id), LR, t);
+        self.wq.adam(&grad_of(&storage, q_id), LR, t);
+        self.wk.adam(&grad_of(&storage, k_id), LR, t);
+        self.wv.adam(&grad_of(&storage, v_id), LR, t);
+        self.w_mem.adam(&grad_of(&storage, mem_id), LR, t);
+        self.w_dec.adam(&grad_of(&storage, dec_id), LR, t);
+
+        loss_val
+    }
+
+    /// Run the forward pass off-tape (no grad) and recover `M` + prediction.
+    fn run(&self, window: &[f32], l: usize) -> Forward {
+        let din = self.cfg.d_in;
+        assert_eq!(window.len(), l * din, "window must be [l, d_in]");
+        // Non-grad leaves: we only read the node values back, so the (harmless)
+        // grad-tracking is never used and nothing calls backward here.
+        let we = Ad::from_data(&self.w_embed.data, self.w_embed.shape(), &self.dev);
+        let q_w = Ad::from_data(&self.wq.data, self.wq.shape(), &self.dev);
+        let k_w = Ad::from_data(&self.wk.data, self.wk.shape(), &self.dev);
+        let v_w = Ad::from_data(&self.wv.data, self.wv.shape(), &self.dev);
+        let mem_w = Ad::from_data(&self.w_mem.data, self.w_mem.shape(), &self.dev);
+        let dec_w = Ad::from_data(&self.w_dec.data, self.w_dec.shape(), &self.dev);
+        let (memory, prediction) = self.graph(window, l, &we, &q_w, &k_w, &v_w, &mem_w, &dec_w);
+        Forward {
+            memory: Ad::to_data(&memory),
+            prediction: Ad::to_data(&prediction),
+        }
+    }
+
+    /// Shared tape graph: build the attention set-encoder, return the scalar SSE
+    /// loss against `target`.
+    #[allow(clippy::too_many_arguments)]
+    fn forward_graph(
+        &self,
+        window: &[f32],
+        l: usize,
+        target: &[f32],
+        we: &<Ad as Backend>::TensorPrimitive<2>,
+        q_w: &<Ad as Backend>::TensorPrimitive<2>,
+        k_w: &<Ad as Backend>::TensorPrimitive<2>,
+        v_w: &<Ad as Backend>::TensorPrimitive<2>,
+        mem_w: &<Ad as Backend>::TensorPrimitive<2>,
+        dec_w: &<Ad as Backend>::TensorPrimitive<2>,
+    ) -> <Ad as Backend>::TensorPrimitive<1> {
+        let din = self.cfg.d_in;
+        let (_memory, pred) = self.graph(window, l, we, q_w, k_w, v_w, mem_w, dec_w);
+        let target_t = Ad::from_data(target, [1, din], &self.dev);
+        let err = Ad::sub(pred, target_t);
+        Ad::sum(Ad::mul(err.clone(), err))
+    }
+
+    /// Build the attention set-encoder on the tape, returning the fixed-size
+    /// memory node `M` ([1, d_mem]) and the prediction node `ŷ` ([1, d_in]).
+    /// Both the training graph and off-tape inference share this one builder, so
+    /// there is a single source of truth for the architecture.
+    #[allow(clippy::too_many_arguments)]
+    fn graph(
+        &self,
+        window: &[f32],
+        l: usize,
+        we: &<Ad as Backend>::TensorPrimitive<2>,
+        q_w: &<Ad as Backend>::TensorPrimitive<2>,
+        k_w: &<Ad as Backend>::TensorPrimitive<2>,
+        v_w: &<Ad as Backend>::TensorPrimitive<2>,
+        mem_w: &<Ad as Backend>::TensorPrimitive<2>,
+        dec_w: &<Ad as Backend>::TensorPrimitive<2>,
+    ) -> (
+        <Ad as Backend>::TensorPrimitive<2>,
+        <Ad as Backend>::TensorPrimitive<2>,
+    ) {
+        let din = self.cfg.d_in;
+        let x = Ad::from_data(window, [l, din], &self.dev);
+        let e = Ad::matmul(x, we.clone()); // [l, d_model]
+        // e fans out to q/k/v and the residual (four uses).
+        let q = Ad::matmul(e.clone(), q_w.clone());
+        let k = Ad::matmul(e.clone(), k_w.clone());
+        let v = Ad::matmul(e.clone(), v_w.clone());
+        let scores = Ad::matmul(q, Ad::transpose(k)); // [l, l]
+        let attn = Ad::softmax(scores, 1);
+        let ctx = Ad::matmul(attn, v); // [l, d_model]
+        let h = Ad::add(e, ctx); // residual [l, d_model]
+        // Mean-pool over the L tokens → [1, d_model] (a set-symmetric reduction).
+        let mean_row = vec![1.0f32 / l as f32; l];
+        let pool = Ad::matmul(Ad::from_data(&mean_row, [1, l], &self.dev), h);
+        let memory = Ad::matmul(pool, mem_w.clone()); // [1, d_mem]
+        let prediction = Ad::matmul(memory.clone(), dec_w.clone()); // [1, d_in]
+        (memory, prediction)
+    }
+
+    /// Canonical `(name, flat-weights, [rows, cols])` view of every parameter,
+    /// in a stable order. The caller (omni-think facade) serializes these to
+    /// safetensors and seals them with BLAKE3.
+    #[must_use]
+    pub fn named_params(&self) -> Vec<(&'static str, Vec<f32>, [usize; 2])> {
+        vec![
+            ("w_embed", self.w_embed.data.clone(), self.w_embed.shape()),
+            ("wq", self.wq.data.clone(), self.wq.shape()),
+            ("wk", self.wk.data.clone(), self.wk.shape()),
+            ("wv", self.wv.data.clone(), self.wv.shape()),
+            ("w_mem", self.w_mem.data.clone(), self.w_mem.shape()),
+            ("w_dec", self.w_dec.data.clone(), self.w_dec.shape()),
+        ]
+    }
+
+    /// Overwrite parameters from a `(name → flat-weights)` lookup, e.g. after
+    /// loading a safetensors checkpoint. Adam moments are reset (the optimizer
+    /// state is not persisted).
+    ///
+    /// # Errors
+    /// Returns `Err(name)` if a parameter is missing or has the wrong length.
+    pub fn set_named_params(
+        &mut self,
+        get: impl Fn(&str) -> Option<Vec<f32>>,
+    ) -> Result<(), String> {
+        for (name, target) in [
+            ("w_embed", &mut self.w_embed),
+            ("wq", &mut self.wq),
+            ("wk", &mut self.wk),
+            ("wv", &mut self.wv),
+            ("w_mem", &mut self.w_mem),
+            ("w_dec", &mut self.w_dec),
+        ] {
+            let want = target.data.len();
+            let got = get(name).ok_or_else(|| format!("missing param `{name}`"))?;
+            if got.len() != want {
+                return Err(format!(
+                    "param `{name}` length {} != expected {want}",
+                    got.len(),
+                ));
+            }
+            target.data = got;
+            target.m = vec![0.0; want];
+            target.v = vec![0.0; want];
+        }
+        self.step = 0;
+        Ok(())
+    }
+}
+
+/// Read a 2-D leaf gradient out of the storage as a flat `Vec`.
+fn grad_of(storage: &rtx_autograd::autodiff::GradientStorage<CpuBackend>, id: usize) -> Vec<f32> {
+    match storage
+        .get(TensorId(id))
+        .unwrap_or_else(|| panic!("no gradient for leaf {id}"))
+    {
+        GradTensor::D2(t) => t.to_vec(),
+        other => panic!("expected D2 gradient, got {other:?}"),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Build a fixed window whose target is the per-dim mean of the window —
+    /// a task that *requires* aggregating across time (memory), so a teacher
+    /// that learns it has built a genuine predictive state.
+    fn window_and_mean_target(l: usize, d_in: usize, seed: u64) -> (Vec<f32>, Vec<f32>) {
+        let x = lcg_init(l * d_in, 1.0, seed);
+        let mut target = vec![0.0f32; d_in];
+        for d in 0..d_in {
+            let mut acc = 0.0;
+            for t in 0..l {
+                acc += x[t * d_in + d];
+            }
+            target[d] = acc / l as f32;
+        }
+        (x, target)
+    }
+
+    #[test]
+    fn teacher_trains_and_loss_decreases() {
+        let cfg = SetEncoderConfig::new(4, 8, 6, 42);
+        let mut teacher = SetEncoderTeacher::new(cfg);
+        let (x, target) = window_and_mean_target(6, 4, 7);
+
+        let first = teacher.train_step(&x, 6, &target);
+        let mut last = first;
+        for _ in 0..400 {
+            last = teacher.train_step(&x, 6, &target);
+            assert!(last.is_finite(), "non-finite loss");
+        }
+        assert!(
+            last < first * 0.2,
+            "teacher did not learn: {first:.5} -> {last:.5}"
+        );
+    }
+
+    #[test]
+    fn encode_is_deterministic_and_fixed_size() {
+        let cfg = SetEncoderConfig::new(4, 8, 6, 1);
+        let teacher = SetEncoderTeacher::new(cfg);
+        let (x, _) = window_and_mean_target(5, 4, 3);
+        let m1 = teacher.encode(&x, 5);
+        let m2 = teacher.encode(&x, 5);
+        assert_eq!(m1.len(), 6, "memory must be fixed size d_mem");
+        assert_eq!(m1, m2, "encode must be deterministic");
+        assert!(m1.iter().all(|v| v.is_finite()));
+    }
+
+    #[test]
+    fn params_round_trip() {
+        let cfg = SetEncoderConfig::new(3, 5, 4, 9);
+        let mut a = SetEncoderTeacher::new(cfg);
+        let (x, target) = window_and_mean_target(4, 3, 2);
+        for _ in 0..20 {
+            a.train_step(&x, 4, &target);
+        }
+        let saved = a.named_params();
+        let mut b = SetEncoderTeacher::new(cfg);
+        b.set_named_params(|name| {
+            saved
+                .iter()
+                .find(|(n, _, _)| *n == name)
+                .map(|(_, w, _)| w.clone())
+        })
+        .expect("round-trip");
+        assert_eq!(a.encode(&x, 4), b.encode(&x, 4), "weights must round-trip");
+    }
+}

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

@@ -0,0 +1,126 @@
+//! D309 equivalence pin — the single-step `MambaRecurrence` is numerically
+//! identical to the full-sequence `MambaBlock::forward`.
+//!
+//! This is the precondition the SMT plan flags as the highest remaining risk:
+//! before a recurrent updater can be behaviorally cloned, its exposed
+//! `(state, x_t) → (state', y_t)` step MUST reproduce the forward exactly.
+//! We run a window through `forward`, then through `step` one timestep at a
+//! time from a zero state, and require the outputs to match to the f32 floor.
+//!
+//! The block uses the `active_block` trick (wide `Δ`) so the scan genuinely
+//! contributes — otherwise the near-identity default init would make the test
+//! pass trivially without exercising the recurrence.
+
+use std::collections::HashMap;
+
+use rtx_tensor::{Device, Tensor};
+use rtx_transformers::layers::mamba::{MambaBlock, MambaConfig};
+use rtx_transformers::layers::mamba_step::MambaRecurrence;
+
+const D_MODEL: usize = 8;
+const D_STATE: usize = 16;
+const D_CONV: usize = 4;
+const L: usize = 7;
+
+fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 {
+    a.iter()
+        .zip(b)
+        .map(|(x, y)| (x - y).abs())
+        .fold(0.0, f32::max)
+}
+
+/// A seeded block with `dt_bias` zeroed (`Δ = softplus(0) ≈ 0.69`) so the scan
+/// genuinely drives the output — mirrors `real_selective_scan.rs::active_block`.
+fn active_block(dev: &Device, seed: u64) -> MambaBlock {
+    let cfg = MambaConfig::new(D_MODEL, D_STATE, D_CONV);
+    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, dev, map).expect("rebuild")
+}
+
+/// Deterministic input window `[1, L, D_MODEL]`.
+fn input(dev: &Device) -> (Tensor, Vec<f32>) {
+    let mut s = 0x57E9u64;
+    let data: Vec<f32> = (0..L * D_MODEL)
+        .map(|_| {
+            s = s
+                .wrapping_mul(6_364_136_223_846_793_005)
+                .wrapping_add(1_442_695_040_888_963_407);
+            ((s >> 40) as f32 / (1u64 << 24) as f32) * 2.0 - 1.0
+        })
+        .collect();
+    let t = Tensor::from_vec(data.clone(), &[1, L, D_MODEL], dev).expect("input");
+    (t, data)
+}
+
+#[test]
+fn step_matches_full_forward() {
+    let dev = Device::cpu();
+    let block = active_block(&dev, 0xD309);
+    let (x, flat) = input(&dev);
+
+    // Full-sequence reference.
+    let full = block
+        .forward(&x)
+        .expect("forward")
+        .output
+        .to_vec()
+        .expect("vec");
+    assert_eq!(full.len(), L * D_MODEL);
+
+    // Stepwise: feed one token at a time from a zero state.
+    let rec = MambaRecurrence::from_block(&block).expect("recurrence");
+    assert_eq!(rec.d_model(), D_MODEL);
+    let mut state = rec.init_state();
+    let mut stepwise = Vec::with_capacity(L * D_MODEL);
+    for li in 0..L {
+        let x_t = &flat[li * D_MODEL..(li + 1) * D_MODEL];
+        stepwise.extend(rec.step(x_t, &mut state));
+    }
+
+    let diff = max_abs_diff(&full, &stepwise);
+    eprintln!("D309 step-vs-forward max_abs_diff = {diff:.2e}");
+    assert!(
+        diff < 1e-4,
+        "stepwise recurrence must match the full forward (max_abs_diff {diff:.2e})"
+    );
+
+    // The state must be non-trivial (the scan actually ran).
+    assert_eq!(rec.state_size(), state.hidden().len());
+    assert!(
+        state.hidden().iter().any(|&v| v.abs() > 1e-6),
+        "SSM state should be populated after stepping"
+    );
+}
+
+#[test]
+fn fresh_state_is_zero_and_deterministic() {
+    let dev = Device::cpu();
+    let block = active_block(&dev, 7);
+    let rec = MambaRecurrence::from_block(&block).expect("recurrence");
+    let s = rec.init_state();
+    assert!(s.hidden().iter().all(|&v| v == 0.0), "init state h is zero");
+
+    // Same input from a fresh state → identical output (stateless determinism).
+    let (_, flat) = input(&dev);
+    let run = |rec: &MambaRecurrence| {
+        let mut st = rec.init_state();
+        let mut out = Vec::new();
+        for li in 0..L {
+            out.extend(rec.step(&flat[li * D_MODEL..(li + 1) * D_MODEL], &mut st));
+        }
+        out
+    };
+    assert_eq!(run(&rec), run(&rec), "stepping must be deterministic");
+}