crates/core/rtx-autograd/tests/tape_cpu_gradcheck.rs

@@ -172,3 +172,96 @@ fn gradcheck_softmax_sum() {
     let numerical = fd_grad(&x, loss_of_x);
     assert_close(&analytical, &numerical, "softmax_sum d/dX");
 }
+
+/// `loss = sum(tanh(X))` — tanh VJP (`grad_out * (1 - tanh(x)^2)`). Needed for
+/// the gated memory cell's candidate state.
+#[test]
+fn gradcheck_tanh_sum() {
+    let x = [0.5f32, -1.0, 2.0, 0.25, -0.5, 1.5]; // [2,3]
+
+    let loss_of_x = |xv: &[f32]| -> f32 {
+        let xt = CpuBackend::from_data(xv, [2, 3], &cpu_dev());
+        CpuBackend::to_data(&CpuBackend::sum(CpuBackend::tanh(xt)))[0]
+    };
+
+    let xt = leaf2(&x, [2, 3]);
+    let x_id = xt.id().0;
+    let loss = Ad::sum(Ad::tanh(xt));
+    let storage = backward_impl(
+        &loss,
+        Some(GradTensor::from_d1(CpuBackend::ones([1], &cpu_dev()))),
+    )
+    .expect("backward");
+
+    let analytical = grad2(&storage, x_id);
+    let numerical = fd_grad(&x, loss_of_x);
+    assert_close(&analytical, &numerical, "tanh_sum d/dX");
+}
+
+/// `loss = sum(sigmoid(X))` — sigmoid VJP (`grad_out * σ(x) * (1 - σ(x))`).
+/// Needed for the gated memory cell's update gate.
+#[test]
+fn gradcheck_sigmoid_sum() {
+    let x = [0.5f32, -1.0, 2.0, 0.25, -0.5, 1.5]; // [2,3]
+
+    let loss_of_x = |xv: &[f32]| -> f32 {
+        let xt = CpuBackend::from_data(xv, [2, 3], &cpu_dev());
+        CpuBackend::to_data(&CpuBackend::sum(CpuBackend::sigmoid(xt)))[0]
+    };
+
+    let xt = leaf2(&x, [2, 3]);
+    let x_id = xt.id().0;
+    let loss = Ad::sum(Ad::sigmoid(xt));
+    let storage = backward_impl(
+        &loss,
+        Some(GradTensor::from_d1(CpuBackend::ones([1], &cpu_dev()))),
+    )
+    .expect("backward");
+
+    let analytical = grad2(&storage, x_id);
+    let numerical = fd_grad(&x, loss_of_x);
+    assert_close(&analytical, &numerical, "sigmoid_sum d/dX");
+}
+
+/// The exact composite the gated memory cell uses:
+/// `loss = sum((1 - z) ⊙ m + z ⊙ c)` with `z = sigmoid(A)`, `c = tanh(B)`, `m`
+/// constant — checks sigmoid·tanh·sub·mul·add compose with correct gradients
+/// (vary the gate logits `A`).
+#[test]
+fn gradcheck_gated_update() {
+    let a = [0.3f32, -0.7, 1.2, 0.1]; // gate logits [2,2]
+    let bb = [0.5f32, -1.0, 0.25, 2.0]; // candidate logits [2,2]
+    let m = [0.4f32, -0.2, 0.9, -0.6]; // previous memory (constant) [2,2]
+
+    let loss_of_a = |av: &[f32]| -> f32 {
+        let at = CpuBackend::from_data(av, [2, 2], &cpu_dev());
+        let bt = CpuBackend::from_data(&bb, [2, 2], &cpu_dev());
+        let mt = CpuBackend::from_data(&m, [2, 2], &cpu_dev());
+        let z = CpuBackend::sigmoid(at);
+        let c = CpuBackend::tanh(bt);
+        let one = CpuBackend::ones([2, 2], &cpu_dev());
+        let keep = CpuBackend::sub(one, z.clone());
+        let new = CpuBackend::add(CpuBackend::mul(keep, mt), CpuBackend::mul(z, c));
+        CpuBackend::to_data(&CpuBackend::sum(new))[0]
+    };
+
+    let at = leaf2(&a, [2, 2]);
+    let bt = leaf2(&bb, [2, 2]);
+    let mt = Ad::from_data(&m, [2, 2], &ad_dev());
+    let a_id = at.id().0;
+    let z = Ad::sigmoid(at);
+    let c = Ad::tanh(bt);
+    let one = Ad::from_data(&[1.0f32; 4], [2, 2], &ad_dev());
+    let keep = Ad::sub(one, z.clone());
+    let new = Ad::add(Ad::mul(keep, mt), Ad::mul(z, c));
+    let loss = Ad::sum(new);
+    let storage = backward_impl(
+        &loss,
+        Some(GradTensor::from_d1(CpuBackend::ones([1], &cpu_dev()))),
+    )
+    .expect("backward");
+
+    let analytical = grad2(&storage, a_id);
+    let numerical = fd_grad(&a, loss_of_a);
+    assert_close(&analytical, &numerical, "gated_update d/d(gate logits)");
+}

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

@@ -0,0 +1,372 @@
+//! `GatedMemoryUpdater` — a GRU-style **gated** recurrent memory cell for
+//! Supervised Memory Training (Phase F / D314).
+//!
+//! The residual-MLP [`ClonedMemoryUpdater`](super::cloned_memory_updater) cell
+//! (`M_t = ρ·M_{t-1} + Δ` + clamp) cannot both *jump* at a regime switch and
+//! *average* in steady state, so when streamed it lags through non-stationary
+//! switches. This cell replaces the fixed leaky decay with a learned per-dim
+//! **update gate**, so the recurrence adapts step by step:
+//! ```text
+//!   z   = sigmoid([M_{t-1} ‖ x] · W_z)     # update gate ∈ (0,1), per memory dim
+//!   c   = tanh([M_{t-1} ‖ x] · W_c)        # candidate state ∈ (-1,1)
+//!   M_t = (1 - z) ⊙ M_{t-1} + z ⊙ c        # z≈1 → jump; z≈0 → hold + denoise
+//!   x̂   = M_t · W_read                     # predict-the-future readout
+//! ```
+//! Because `M_t` is a **convex combination** of `M_{t-1}` and the bounded
+//! candidate `c ∈ (-1,1)`, `|M_t| ≤ max(|M_0|, 1)` for all `t` — the recurrence
+//! is a non-expansion, so a free autoregressive rollout stays bounded **without**
+//! a clamp.
+//!
+//! Trains on `Autodiff<CpuBackend>` with a self-contained deterministic Adam,
+//! exposing the same method surface as `ClonedMemoryUpdater` (so the omni-think
+//! facade is cell-generic). sigmoid/tanh/mul/add/sub VJPs are finite-difference
+//! gated in `rtx-autograd/tests/tape_cpu_gradcheck.rs`.
+
+use rtx_autograd::autodiff::{Autodiff, AutodiffDevice, GradTensor, TensorId, backward_impl};
+use rtx_backend::{AutodiffBackend, Backend};
+use rtx_backend_cpu::CpuBackend;
+
+use super::cloned_memory_updater::ClonedUpdaterConfig;
+
+type Ad = Autodiff<CpuBackend>;
+
+/// 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()
+}
+
+/// Numerically-naive sigmoid matching `rtx-backend-cpu`'s, so off-tape `step`
+/// agrees with the trained graph.
+#[inline]
+fn sigmoid(x: f32) -> f32 {
+    1.0 / (1.0 + (-x).exp())
+}
+
+/// 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:?}"),
+    }
+}
+
+/// A GRU-style gated recurrent memory updater + predict-the-future readout.
+pub struct GatedMemoryUpdater {
+    cfg: ClonedUpdaterConfig,
+    w_z: Param,    // [d_mem + d_in, d_mem]  update gate
+    w_c: Param,    // [d_mem + d_in, d_mem]  candidate
+    w_read: Param, // [d_mem, d_in]
+    step: u32,
+    dev: AutodiffDevice<CpuBackend>,
+}
+
+impl GatedMemoryUpdater {
+    /// Build with deterministic, seeded weights. `cfg.d_hidden` is unused by the
+    /// gated cell (gates project `[M ‖ x]` straight to `d_mem`); it is kept in
+    /// the shared config so this is a drop-in for `ClonedMemoryUpdater`.
+    #[must_use]
+    pub fn new(cfg: ClonedUpdaterConfig) -> Self {
+        let s = |i: u64| cfg.seed.wrapping_mul(0x100_0000_01b3).wrapping_add(i);
+        let fan = cfg.d_mem + cfg.d_in;
+        let sc = 1.0 / (fan as f32).sqrt();
+        Self {
+            w_z: Param::new(fan, cfg.d_mem, sc, s(1)),
+            w_c: Param::new(fan, cfg.d_mem, sc, s(2)),
+            w_read: Param::new(cfg.d_mem, cfg.d_in, 1.0 / (cfg.d_mem as f32).sqrt(), s(3)),
+            step: 0,
+            dev: AutodiffDevice::<CpuBackend>::default(),
+            cfg,
+        }
+    }
+
+    /// Memory size `|M|`.
+    #[must_use]
+    pub const fn d_mem(&self) -> usize {
+        self.cfg.d_mem
+    }
+
+    /// Off-tape gated step `g(prev_mem, x) → new_mem` (`[d_mem]`).
+    ///
+    /// # Panics
+    /// Panics on shape mismatch.
+    #[must_use]
+    pub fn step(&self, prev_mem: &[f32], x: &[f32]) -> Vec<f32> {
+        let dm = self.cfg.d_mem;
+        let din = self.cfg.d_in;
+        assert_eq!(prev_mem.len(), dm);
+        assert_eq!(x.len(), din);
+        let cat_at = |i: usize| if i < dm { prev_mem[i] } else { x[i - dm] };
+        let fan = dm + din;
+        let mut out = vec![0.0f32; dm];
+        for j in 0..dm {
+            let mut zlin = 0.0f32;
+            let mut clin = 0.0f32;
+            for i in 0..fan {
+                let ci = cat_at(i);
+                zlin += ci * self.w_z.data[i * dm + j];
+                clin += ci * self.w_c.data[i * dm + j];
+            }
+            let z = sigmoid(zlin);
+            let c = clin.tanh();
+            out[j] = (1.0 - z) * prev_mem[j] + z * c;
+        }
+        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> {
+        let dm = self.cfg.d_mem;
+        let din = self.cfg.d_in;
+        assert_eq!(mem.len(), dm);
+        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
+    }
+
+    /// Build the gated update graph on the tape, returning the `new_mem`
+    /// ([1, d_mem]) node. Shared by both training entry points.
+    fn graph(
+        &self,
+        prev_mem: &[f32],
+        x: &[f32],
+        w_z: &<Ad as Backend>::TensorPrimitive<2>,
+        w_c: &<Ad as Backend>::TensorPrimitive<2>,
+    ) -> <Ad as Backend>::TensorPrimitive<2> {
+        let dm = self.cfg.d_mem;
+        let din = self.cfg.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 z = Ad::sigmoid(Ad::matmul(cat_t.clone(), w_z.clone())); // [1, d_mem]
+        let c = Ad::tanh(Ad::matmul(cat_t, w_c.clone()));
+        let one = Ad::from_data(&vec![1.0f32; dm], [1, dm], &self.dev);
+        let keep = Ad::sub(one, z.clone());
+        // M_t = (1 - z) ⊙ prev + z ⊙ c
+        Ad::add(Ad::mul(keep, prev_t), Ad::mul(z, c))
+    }
+
+    /// One behavioral-cloning step: train the gate + candidate to map
+    /// `(prev_mem, x) → target_mem` and the readout to predict `target_next`.
+    /// Returns the total loss before the update.
+    ///
+    /// # Panics
+    /// Panics on 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_z, z_id) = self.w_z.leaf(&self.dev);
+        let (w_c, c_id) = self.w_c.leaf(&self.dev);
+        let (w_read, read_id) = self.w_read.leaf(&self.dev);
+
+        let new_mem = self.graph(prev_mem, x, &w_z, &w_c);
+        let pred = Ad::matmul(new_mem.clone(), w_read);
+        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(&loss);
+        if !loss_val.is_finite() {
+            return loss_val;
+        }
+        self.step += 1;
+        let t = self.step;
+        const LR: f32 = 0.01;
+        self.w_z.adam(&clip(grad_of(&storage, z_id)), LR, t);
+        self.w_c.adam(&clip(grad_of(&storage, c_id)), LR, t);
+        self.w_read.adam(&clip(grad_of(&storage, read_id)), LR, t);
+        loss_val
+    }
+
+    /// DAgger memory-only correction at learning rate `lr`: train just the gate +
+    /// candidate to map `(prev_mem, x) → target_mem`, leaving the readout intact.
+    ///
+    /// # Panics
+    /// Panics on shape mismatch.
+    pub fn train_step_memory(
+        &mut self,
+        prev_mem: &[f32],
+        x: &[f32],
+        target_mem: &[f32],
+        lr: f32,
+    ) -> f32 {
+        let dm = self.cfg.d_mem;
+        assert_eq!(prev_mem.len(), dm);
+        assert_eq!(x.len(), self.cfg.d_in);
+        assert_eq!(target_mem.len(), dm);
+
+        let (w_z, z_id) = self.w_z.leaf(&self.dev);
+        let (w_c, c_id) = self.w_c.leaf(&self.dev);
+        let new_mem = self.graph(prev_mem, x, &w_z, &w_c);
+        let tgt_mem_t = Ad::from_data(target_mem, [1, dm], &self.dev);
+        let mem_err = Ad::sub(new_mem, tgt_mem_t);
+        let loss = Ad::sum(Ad::mul(mem_err.clone(), mem_err));
+        let loss_val = Ad::to_data(&loss)[0];
+
+        let storage = backward(&loss);
+        if !loss_val.is_finite() {
+            return loss_val;
+        }
+        self.step += 1;
+        let t = self.step;
+        self.w_z.adam(&clip(grad_of(&storage, z_id)), lr, t);
+        self.w_c.adam(&clip(grad_of(&storage, c_id)), lr, t);
+        loss_val
+    }
+}
+
+/// Run the tape backward from a scalar loss.
+fn backward(
+    loss: &<Ad as Backend>::TensorPrimitive<1>,
+) -> rtx_autograd::autodiff::GradientStorage<CpuBackend> {
+    backward_impl(
+        loss,
+        Some(GradTensor::from_d1(CpuBackend::ones(
+            [1],
+            &<CpuBackend as Backend>::Device::default(),
+        ))),
+    )
+    .expect("backward")
+}
+
+/// Per-element gradient clip for stability on extreme rollout states.
+fn clip(mut g: Vec<f32>) -> Vec<f32> {
+    const CLIP: f32 = 1.0;
+    for v in &mut g {
+        *v = v.clamp(-CLIP, CLIP);
+    }
+    g
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn gated_updater_learns_a_transition() {
+        let cfg = ClonedUpdaterConfig::new(4, 6, 16, 1);
+        let mut g = GatedMemoryUpdater::new(cfg);
+        let prev = lcg_init(6, 0.5, 10);
+        let x = lcg_init(4, 1.0, 11);
+        let target_mem = lcg_init(6, 0.4, 12);
+        let target_next = lcg_init(4, 0.4, 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.5,
+            "gated updater did not learn: {first:.4} -> {last:.4}"
+        );
+    }
+
+    #[test]
+    fn free_rollout_is_bounded_without_clamp() {
+        // The convex update is a non-expansion: |M_t| ≤ max(|M_0|, 1) forever.
+        let cfg = ClonedUpdaterConfig::new(4, 6, 8, 2);
+        let g = GatedMemoryUpdater::new(cfg);
+        let mut mem = vec![0.3f32; 6];
+        for t in 0..500 {
+            let x = lcg_init(4, 1.0, t as u64);
+            mem = g.step(&mem, &x);
+            assert!(
+                mem.iter().all(|v| v.abs() <= 1.0 + 1e-4),
+                "memory left [-1,1] at step {t}"
+            );
+        }
+    }
+}

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

@@ -66,6 +66,10 @@ pub mod set_encoder_teacher;
 // trajectory, no BPTT — see `cloned_memory_updater`).
 pub mod cloned_memory_updater;
 
+// SMT gated (GRU-style) memory updater — switch-tracking, bounded recurrence
+// for non-stationary streams (see `gated_memory_updater`).
+pub mod gated_memory_updater;
+
 // Ring attention for long context
 pub mod ring_attention;
 // pub mod mamba_integration;