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

@@ -386,6 +386,106 @@ impl ClonedMemoryUpdater {
         self.w_mem.adam(&clip(grad_of(&storage, mem_id)), lr, t);
         loss_val
     }
+
+    /// **Multi-step rollout training (truncated K-step BPTT).** Starting from
+    /// `init_mem`, roll the cell forward over `xs` **on its own memory** (not
+    /// teacher-forced), accumulate the predict-the-future loss `Σ_k ‖readout(M_k)
+    /// − targets[k]‖²`, and backprop through the *whole* K-step graph — so the
+    /// cell is optimized for the free-rollout behavior it is actually evaluated
+    /// on, directly attacking the autoregressive drift that one-step BC leaves.
+    ///
+    /// To keep the K-step graph within the gradient-checked op set (matmul / gelu
+    /// / add / mul / sub / sum — no tensor concat), `W_in` is split into its
+    /// memory rows (applied to the recurrent state `M`) and its input rows
+    /// (applied to `x`); the two gradient halves are re-stacked for the Adam
+    /// update. Returns the mean per-step loss.
+    ///
+    /// # Panics
+    /// Panics if `xs.len() != targets.len()`, or on any shape mismatch.
+    pub fn train_rollout(
+        &mut self,
+        init_mem: &[f32],
+        xs: &[Vec<f32>],
+        targets: &[Vec<f32>],
+        lr: f32,
+    ) -> f32 {
+        let dm = self.cfg.d_mem;
+        let din = self.cfg.d_in;
+        let dh = self.cfg.d_hidden;
+        assert_eq!(init_mem.len(), dm);
+        assert_eq!(xs.len(), targets.len());
+        assert!(!xs.is_empty());
+
+        // W_in split into memory rows [0..dm] and input rows [dm..dm+din].
+        let w_in_m = Ad::require_grad(Ad::from_data(
+            &self.w_in.data[..dm * dh],
+            [dm, dh],
+            &self.dev,
+        ));
+        let w_in_x = Ad::require_grad(Ad::from_data(
+            &self.w_in.data[dm * dh..],
+            [din, dh],
+            &self.dev,
+        ));
+        let (w_mem, mem_id) = self.w_mem.leaf(&self.dev);
+        let (w_read, read_id) = self.w_read.leaf(&self.dev);
+        let in_m_id = w_in_m.id().0;
+        let in_x_id = w_in_x.id().0;
+        let decay_t = Ad::from_data(&vec![DECAY; dm], [1, dm], &self.dev);
+
+        let mut m = Ad::from_data(init_mem, [1, dm], &self.dev); // rollout start (const)
+        let mut loss: Option<<Ad as Backend>::TensorPrimitive<1>> = None;
+        for (x, target) in xs.iter().zip(targets.iter()) {
+            let x_t = Ad::from_data(x, [1, din], &self.dev);
+            let preact = Ad::add(
+                Ad::matmul(m.clone(), w_in_m.clone()),
+                Ad::matmul(x_t, w_in_x.clone()),
+            );
+            let h = Ad::gelu(preact);
+            let delta = Ad::matmul(h, w_mem.clone());
+            m = Ad::add(Ad::mul(m, decay_t.clone()), delta); // leaky update (no clamp; K short)
+            let pred = Ad::matmul(m.clone(), w_read.clone());
+            let err = Ad::sub(pred, Ad::from_data(target, [1, din], &self.dev));
+            let step_loss = Ad::sum(Ad::mul(err.clone(), err));
+            loss = Some(match loss {
+                None => step_loss,
+                Some(acc) => Ad::add(acc, step_loss),
+            });
+        }
+        let loss = loss.expect("non-empty rollout");
+        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");
+
+        let k = xs.len() as f32;
+        if !loss_val.is_finite() {
+            return loss_val / k;
+        }
+        self.step += 1;
+        let t = self.step;
+        const CLIP: f32 = 1.0;
+        let clip = |mut g: Vec<f32>| {
+            for v in &mut g {
+                *v = v.clamp(-CLIP, CLIP);
+            }
+            g
+        };
+        // Re-stack the two W_in gradient halves (memory rows then input rows) into
+        // the full row-major W_in gradient.
+        let mut grad_in = grad_of(&storage, in_m_id);
+        grad_in.extend(grad_of(&storage, in_x_id));
+        self.w_in.adam(&clip(grad_in), lr, t);
+        self.w_mem.adam(&clip(grad_of(&storage, mem_id)), lr, t);
+        self.w_read.adam(&clip(grad_of(&storage, read_id)), lr, t);
+        loss_val / k
+    }
 }
 
 /// Exact-ish GELU matching the tape's `Ad::gelu` (tanh approximation), so the
@@ -432,4 +532,25 @@ mod tests {
         assert_eq!(g.predict(&new_mem).len(), 4);
         assert!(new_mem.iter().all(|v| v.is_finite()));
     }
+
+    #[test]
+    fn rollout_training_reduces_loss() {
+        // Truncated K-step BPTT drives down the rollout's predict-the-future loss.
+        let cfg = ClonedUpdaterConfig::new(4, 6, 16, 1);
+        let mut g = ClonedMemoryUpdater::new(cfg);
+        let init = lcg_init(6, 0.3, 20);
+        let xs: Vec<Vec<f32>> = (0..6).map(|k| lcg_init(4, 1.0, 30 + k)).collect();
+        let targets: Vec<Vec<f32>> = (0..6).map(|k| lcg_init(4, 0.4, 60 + k)).collect();
+
+        let first = g.train_rollout(&init, &xs, &targets, 0.01);
+        let mut last = first;
+        for _ in 0..400 {
+            last = g.train_rollout(&init, &xs, &targets, 0.01);
+            assert!(last.is_finite());
+        }
+        assert!(
+            last < first * 0.5,
+            "rollout training did not reduce loss: {first:.4} -> {last:.4}"
+        );
+    }
 }