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

@@ -0,0 +1,333 @@
+//! `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]   (residual update)
+//!   x̂_{t+1} = M_t · W_read                   [d_in]    (predict-the-future head)
+//! ```
+//! The residual update (`M_t = M_{t-1} + Δ`) makes the cell learn the *increment*
+//! the new token induces — natural for a recurrent state and the key to
+//! extrapolating past the teacher's training horizon (the recurrence rule is
+//! length-invariant).
+//!
+//! 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>;
+
+/// 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
+        let mut out = prev_mem.to_vec();
+        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];
+            }
+        }
+        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 new_mem = Ad::add(prev_t, delta); // residual 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/mod.rs

@@ -62,6 +62,10 @@ pub mod metal_mamba;
 // 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;