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

@@ -15,13 +15,13 @@
 //! ```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)
+//!   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 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).
+//! 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
@@ -34,6 +34,16 @@ 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 {
@@ -191,13 +201,20 @@ impl ClonedMemoryUpdater {
             }
             *hk = gelu_scalar(acc);
         }
-        // new_mem = prev_mem + h · W_mem
-        let mut out = prev_mem.to_vec();
+        // 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
     }
 
@@ -254,7 +271,8 @@ impl ClonedMemoryUpdater {
 
         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 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);