SMT E1: SetEncoderTeacher — the predictive-state oracle (time-parallel, tape-trained)

Promotes the proven attention set-encoder smoke test (PR #4,
tape_train_smoke.rs) into a reusable rtx-transformers layer for Supervised
Memory Training. The teacher maps a window of past tokens to a fixed-size
memory M via embed → self-attention → residual → mean-pool → memory
projection, with a decoder head supervised by predict-the-future MSE so that
M becomes a sufficient statistic of the past.

- Trains end-to-end on Autodiff<CpuBackend> (the gradient-correct real backend
  from PR #3), with a self-contained deterministic host-side Adam.
- Time-parallel by construction (one window → one memory, no recurrence to
  unroll) — this is the oracle whose trajectory the recurrent Mamba updater is
  later behaviorally cloned against, so the recurrent net never needs BPTT.
- Exposes named_params/set_named_params so the caller (omni-think's
  PredictiveStateTeacher facade) owns safetensors persistence + BLAKE3 sealing.
- Adds rtx-backend + rtx-backend-cpu deps (the tape needs a concrete backend).

Tests: teacher trains (loss >5x drop), encode is deterministic + fixed-size,
params round-trip. fmt + clippy(-D warnings) clean on the new module.

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

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/mod.rs

@@ -57,6 +57,10 @@ pub mod mixture_of_experts;
 pub mod mamba;
 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;
+
 // 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");
+    }
+}