SMT D319 (rustytorch): learned-[MEM]-query (content-addressable) teacher pool

Adds an opt-in learned-query attention pool to SetEncoderTeacher
(SetEncoderConfig::with_learned_pool): pool = softmax(q_mem·hᵀ)·h instead of the
fixed mean/decay pool — Isola's transformer-teacher [MEM] query. New q_mem param
(threaded through graph/train_step/run/named_params; only updated in learned-pool
mode). Default off → existing teachers byte-unchanged.

Finding (test teacher_content_addresses_selective_retrieval): on a selective-
retrieval task (signal in one marked token among distractors) BOTH the mean-pool
and learned-query teachers recover the marked token to low MSE (~0.0006 / ~0.002)
— because the self-attention layer already routes the marked token's signal to
every position before the pool. So the mean-pool was NOT the recall bottleneck
(correcting the D318 hypothesis): the teacher can content-address; the real
recall bottleneck is the recurrent *cell* that imitates it. The learned-query
pool is shipped as an equally-capable, faithful-to-the-talk alternative.

All 5 teacher unit tests pass; clippy(-D)/fmt clean.

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

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

@@ -71,6 +71,14 @@ pub struct SetEncoderConfig {
     /// *sharper* (concentrates on the most recent tokens → fast switch reaction,
     /// less denoising); larger = wider/smoother window. Default 0.85.
     pub recency_decay: f32,
+    /// **Content-addressable pooling.** When `false` (default) the window is
+    /// pooled by fixed weights (uniform mean, or recency decay). When `true` a
+    /// **learned `[MEM]` query** attends over the tokens
+    /// (`pool = softmax(q_mem·hᵀ)·h`), so the encoder *learns which tokens to
+    /// read* — the capability a mean/decay pool fundamentally lacks (selective
+    /// retrieval, not just averaging). This is Isola's transformer teacher;
+    /// `false` is the simplified aggregator. Overrides the fixed pool when set.
+    pub learned_pool: bool,
 }
 
 impl SetEncoderConfig {
@@ -84,9 +92,17 @@ impl SetEncoderConfig {
             seed,
             recency: false,
             recency_decay: 0.85,
+            learned_pool: false,
         }
     }
 
+    /// Enable/disable the learned-`[MEM]`-query (content-addressable) pool.
+    #[must_use]
+    pub const fn with_learned_pool(mut self, learned_pool: bool) -> Self {
+        self.learned_pool = learned_pool;
+        self
+    }
+
     /// Enable/disable recency (timestamp-embedded, decay-pooled) mode.
     #[must_use]
     pub const fn with_recency(mut self, recency: bool) -> Self {
@@ -224,6 +240,7 @@ pub struct SetEncoderTeacher {
     wv: Param,      // [d_model, d_model]
     w_mem: Param,   // [d_model, d_mem]
     w_dec: Param,   // [d_mem, d_in]
+    q_mem: Param,   // [1, d_model]  learned [MEM] pooling query (content-addressable)
     /// Global Adam step counter (1-based after the first update).
     step: u32,
     dev: AutodiffDevice<CpuBackend>,
@@ -253,6 +270,7 @@ impl SetEncoderTeacher {
             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)),
+            q_mem: Param::new(1, cfg.d_model, sc(cfg.d_model), s(7)),
             step: 0,
             dev: AutodiffDevice::<CpuBackend>::default(),
             cfg,
@@ -301,8 +319,11 @@ impl SetEncoderTeacher {
         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 (qm, qm_id) = self.q_mem.leaf(&self.dev);
 
-        let loss = self.forward_graph(window, l, target, &we, &q_w, &k_w, &v_w, &mem_w, &dec_w);
+        let loss = self.forward_graph(
+            window, l, target, &we, &q_w, &k_w, &v_w, &mem_w, &dec_w, &qm,
+        );
         let loss_val = Ad::to_data(&loss)[0];
 
         let storage = backward_impl(
@@ -324,6 +345,10 @@ impl SetEncoderTeacher {
         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);
+        // `q_mem` only participates (and receives a gradient) in learned-pool mode.
+        if self.cfg.learned_pool {
+            self.q_mem.adam(&grad_of(&storage, qm_id), LR, t);
+        }
 
         loss_val
     }
@@ -340,7 +365,9 @@ impl SetEncoderTeacher {
         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);
+        let qm = Ad::from_data(&self.q_mem.data, self.q_mem.shape(), &self.dev);
+        let (memory, prediction) =
+            self.graph(window, l, &we, &q_w, &k_w, &v_w, &mem_w, &dec_w, &qm);
         Forward {
             memory: Ad::to_data(&memory),
             prediction: Ad::to_data(&prediction),
@@ -361,9 +388,10 @@ impl SetEncoderTeacher {
         v_w: &<Ad as Backend>::TensorPrimitive<2>,
         mem_w: &<Ad as Backend>::TensorPrimitive<2>,
         dec_w: &<Ad as Backend>::TensorPrimitive<2>,
+        qm: &<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 (_memory, pred) = self.graph(window, l, we, q_w, k_w, v_w, mem_w, dec_w, qm);
         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))
@@ -384,6 +412,7 @@ impl SetEncoderTeacher {
         v_w: &<Ad as Backend>::TensorPrimitive<2>,
         mem_w: &<Ad as Backend>::TensorPrimitive<2>,
         dec_w: &<Ad as Backend>::TensorPrimitive<2>,
+        qm: &<Ad as Backend>::TensorPrimitive<2>,
     ) -> (
         <Ad as Backend>::TensorPrimitive<2>,
         <Ad as Backend>::TensorPrimitive<2>,
@@ -405,11 +434,18 @@ 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]
-        // Pool over the L tokens → [1, d_model]. Stationary mode uses a uniform
+        // Pool the L tokens → [1, d_model]. Learned-query mode lets a learned
-        // mean (order-invariant); recency mode uses an exponential-decay weight
+        // `[MEM]` query attend over the tokens (content-addressable: the encoder
-        // (recent tokens dominate) so the statistic tracks non-stationary state.
+        // *reads* the relevant tokens); otherwise a fixed weight (uniform mean,
+        // or recency-decay) just *averages* them.
+        let pool = if self.cfg.learned_pool {
+            let pool_scores = Ad::matmul(qm.clone(), Ad::transpose(h.clone())); // [1, l]
+            let pool_attn = Ad::softmax(pool_scores, 1);
+            Ad::matmul(pool_attn, h) // [1, d_model]
+        } else {
             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);
+            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)
@@ -427,6 +463,7 @@ impl SetEncoderTeacher {
             ("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()),
+            ("q_mem", self.q_mem.data.clone(), self.q_mem.shape()),
         ]
     }
 
@@ -447,6 +484,7 @@ impl SetEncoderTeacher {
             ("wv", &mut self.wv),
             ("w_mem", &mut self.w_mem),
             ("w_dec", &mut self.w_dec),
+            ("q_mem", &mut self.q_mem),
         ] {
             let want = target.data.len();
             let got = get(name).ok_or_else(|| format!("missing param `{name}`"))?;
@@ -496,6 +534,91 @@ mod tests {
         (x, target)
     }
 
+    /// D319 — selective-retrieval task: a window of `L` tokens where exactly one
+    /// "marked" token (dim 0 high) carries the signal in dims 1.. and all others
+    /// are distractor noise (dim 0 low). The target is the marked token. A mean
+    /// pool dilutes the signal by `L`; a learned `[MEM]` query can *attend to the
+    /// marker* and read it out. Returns `(window [l*d_in], target [d_in])`.
+    fn selective_sample(l: usize, d_in: usize, seed: u64) -> (Vec<f32>, Vec<f32>) {
+        let marked = (seed as usize) % l;
+        let sig = lcg_init(d_in, 0.7, seed.wrapping_mul(31).wrapping_add(1));
+        let noise = lcg_init(l * d_in, 0.7, seed.wrapping_mul(97).wrapping_add(3));
+        let mut win = vec![0.0f32; l * d_in];
+        let mut target = vec![0.0f32; d_in];
+        for t in 0..l {
+            if t == marked {
+                win[t * d_in] = 1.0; // marker flag
+                for d in 1..d_in {
+                    win[t * d_in + d] = sig[d];
+                }
+            } else {
+                win[t * d_in] = -1.0; // unmarked
+                for d in 1..d_in {
+                    win[t * d_in + d] = noise[t * d_in + d];
+                }
+            }
+        }
+        target[0] = 1.0;
+        for d in 1..d_in {
+            target[d] = sig[d];
+        }
+        (win, target)
+    }
+
+    /// D319 — the teacher **already content-addresses**: on a selective-retrieval
+    /// task (signal in one marked token among distractors) *both* pooling modes
+    /// recover the marked token to low error — because the **self-attention layer
+    /// routes the marked token's signal to every position before the pool**, so a
+    /// mean pool suffices. This refutes the D318 hypothesis that the mean-pool was
+    /// the recall bottleneck: it isn't — the teacher can select. (The learned
+    /// `[MEM]`-query pool, exercised here, is an equally-capable available
+    /// alternative; the genuine recall bottleneck is the recurrent *cell* that
+    /// imitates the teacher, not the teacher.)
+    #[test]
+    fn teacher_content_addresses_selective_retrieval() {
+        const L: usize = 8;
+        const DIN: usize = 4;
+        let train: Vec<_> = (0..20u64)
+            .map(|i| selective_sample(L, DIN, 1000 + i))
+            .collect();
+        let test: Vec<_> = (0..8u64)
+            .map(|i| selective_sample(L, DIN, 9000 + i))
+            .collect();
+
+        let mut mean = SetEncoderTeacher::new(SetEncoderConfig::new(DIN, 16, 8, 7));
+        let mut learned =
+            SetEncoderTeacher::new(SetEncoderConfig::new(DIN, 16, 8, 7).with_learned_pool(true));
+        for _ in 0..250 {
+            for (w, t) in &train {
+                mean.train_step(w, L, t);
+                learned.train_step(w, L, t);
+            }
+        }
+
+        let eval = |teacher: &SetEncoderTeacher| -> f32 {
+            let mut acc = 0.0f32;
+            for (w, t) in &test {
+                let p = teacher.predict(w, L);
+                acc += p.iter().zip(t).map(|(a, b)| (a - b).powi(2)).sum::<f32>() / DIN as f32;
+            }
+            acc / test.len() as f32
+        };
+        let mean_mse = eval(&mean);
+        let learned_mse = eval(&learned);
+        eprintln!(
+            "D319 selective-retrieval MSE  mean-pool={mean_mse:.5}  learned-query={learned_mse:.5}"
+        );
+
+        // Both content-address well (target magnitude ~0.7; a pure averager that
+        // couldn't select would land near ~0.3 MSE). The self-attention does the
+        // selecting; the learned-query pool is a working alternative, not a fix.
+        assert!(
+            mean_mse < 0.02 && learned_mse < 0.02,
+            "teacher should content-address selective retrieval with either pool: \
+             mean {mean_mse:.5}, learned {learned_mse:.5}"
+        );
+    }
+
     #[test]
     fn teacher_trains_and_loss_decreases() {
         let cfg = SetEncoderConfig::new(4, 8, 6, 42);