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

@@ -59,10 +59,18 @@ pub struct SetEncoderConfig {
     /// Deterministic init seed. Two teachers built with the same `(config,
     /// seed)` are bit-identical, so training is reproducible.
     pub seed: u64,
+    /// **Recency / timestamp mode.** When `false` (default) the encoder mean-
+    /// pools the whole window — an order-invariant, *stationary* sufficient
+    /// statistic. When `true` it adds sinusoidal **timestamp embeddings** to the
+    /// token embeddings (so attention can reason about position) and pools with
+    /// an **exponential-decay (recent-weighted)** reduction instead of a uniform
+    /// mean — giving a *recent-window* statistic that tracks non-stationary /
+    /// regime-switching signals (the talk's timestamp-embedded teacher).
+    pub recency: bool,
 }
 
 impl SetEncoderConfig {
-    /// Standard config for a `d_in`-wide latent stream.
+    /// Standard (stationary, mean-pool) 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 {
@@ -70,8 +78,16 @@ impl SetEncoderConfig {
             d_model,
             d_mem,
             seed,
+            recency: false,
         }
     }
+
+    /// Enable/disable recency (timestamp-embedded, decay-pooled) mode.
+    #[must_use]
+    pub const fn with_recency(mut self, recency: bool) -> Self {
+        self.recency = recency;
+        self
+    }
 }
 
 /// One trainable parameter matrix held on the host as a flat row-major `Vec`,
@@ -137,6 +153,38 @@ impl Param {
     }
 }
 
+/// Pooling weights over `l` positions: a uniform mean (stationary), or an
+/// exponential-decay weight `ρ^(l-1-t)` normalized to sum 1 (recency mode), so
+/// the most recent token has weight ∝ 1 and older tokens decay — a recent-window
+/// statistic that tracks non-stationary signals.
+fn pool_weights(l: usize, recency: bool) -> Vec<f32> {
+    if !recency {
+        return vec![1.0f32 / l as f32; l];
+    }
+    const RHO: f32 = 0.85;
+    let mut w: Vec<f32> = (0..l).map(|t| RHO.powi((l - 1 - t) as i32)).collect();
+    let sum: f32 = w.iter().sum();
+    for x in &mut w {
+        *x /= sum;
+    }
+    w
+}
+
+/// Standard sinusoidal positional encoding, row-major `[l, d_model]`. Even dims
+/// use `sin`, odd dims `cos`, with the usual geometric wavelength schedule.
+fn sinusoidal_pos(l: usize, d_model: usize) -> Vec<f32> {
+    let mut out = vec![0.0f32; l * d_model];
+    for t in 0..l {
+        for i in 0..d_model {
+            let pair = (i / 2) as f32;
+            let freq = 1.0f32 / 10_000_f32.powf(2.0 * pair / d_model as f32);
+            let angle = t as f32 * freq;
+            out[t * d_model + i] = if i % 2 == 0 { angle.sin() } else { angle.cos() };
+        }
+    }
+    out
+}
+
 /// 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> {
@@ -327,8 +375,14 @@ impl SetEncoderTeacher {
         <Ad as Backend>::TensorPrimitive<2>,
     ) {
         let din = self.cfg.d_in;
+        let dm = self.cfg.d_model;
         let x = Ad::from_data(window, [l, din], &self.dev);
-        let e = Ad::matmul(x, we.clone()); // [l, d_model]
+        let mut e = Ad::matmul(x, we.clone()); // [l, d_model]
+        // Recency mode: add sinusoidal timestamp embeddings so attention can
+        // reason about position (the talk's reparameterization).
+        if self.cfg.recency {
+            e = Ad::add(e, Ad::from_data(&sinusoidal_pos(l, dm), [l, dm], &self.dev));
+        }
         // 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());
@@ -337,9 +391,11 @@ impl SetEncoderTeacher {
         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);
+        // Pool over the L tokens → [1, d_model]. Stationary mode uses a uniform
+        // mean (order-invariant); recency mode uses an exponential-decay weight
+        // (recent tokens dominate) so the statistic tracks non-stationary state.
+        let pool_row = pool_weights(l, self.cfg.recency);
+        let pool = Ad::matmul(Ad::from_data(&pool_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)
@@ -444,6 +500,25 @@ mod tests {
         );
     }
 
+    #[test]
+    fn recency_mode_trains_and_loss_decreases() {
+        // Timestamp-embedded + decay-pooled teacher still trains end-to-end.
+        let cfg = SetEncoderConfig::new(4, 8, 6, 42).with_recency(true);
+        let mut teacher = SetEncoderTeacher::new(cfg);
+        assert!(teacher.config().recency);
+        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,
+            "recency teacher did not learn: {first:.5} -> {last:.5}"
+        );
+    }
+
     #[test]
     fn encode_is_deterministic_and_fixed_size() {
         let cfg = SetEncoderConfig::new(4, 8, 6, 1);