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

@@ -67,6 +67,10 @@ pub struct SetEncoderConfig {
     /// mean — giving a *recent-window* statistic that tracks non-stationary /
     /// regime-switching signals (the talk's timestamp-embedded teacher).
     pub recency: bool,
+    /// Geometric decay of the recency pool (only used when `recency`). Smaller =
+    /// *sharper* (concentrates on the most recent tokens → fast switch reaction,
+    /// less denoising); larger = wider/smoother window. Default 0.85.
+    pub recency_decay: f32,
 }
 
 impl SetEncoderConfig {
@@ -79,6 +83,7 @@ impl SetEncoderConfig {
             d_mem,
             seed,
             recency: false,
+            recency_decay: 0.85,
         }
     }
 
@@ -88,6 +93,13 @@ impl SetEncoderConfig {
         self.recency = recency;
         self
     }
+
+    /// Set the recency pool's geometric decay (see [`SetEncoderConfig::recency_decay`]).
+    #[must_use]
+    pub const fn with_recency_decay(mut self, decay: f32) -> Self {
+        self.recency_decay = decay;
+        self
+    }
 }
 
 /// One trainable parameter matrix held on the host as a flat row-major `Vec`,
@@ -157,12 +169,14 @@ impl Param {
 /// 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> {
+fn pool_weights(l: usize, recency: bool, rho: f32) -> Vec<f32> {
     if !recency {
         return vec![1.0f32 / l as f32; l];
     }
-    const RHO: f32 = 0.85;
+    // Geometric decay: a smaller `rho` concentrates on the most recent tokens (a
-    let mut w: Vec<f32> = (0..l).map(|t| RHO.powi((l - 1 - t) as i32)).collect();
+    // *sharper* oracle that reacts fast to regime switches but denoises less); a
+    // larger `rho` is a wider, smoother recent window.
+    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;
@@ -394,7 +408,7 @@ impl SetEncoderTeacher {
         // 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_row = pool_weights(l, self.cfg.recency, self.cfg.recency_decay);
         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]