rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

SMT D319: learned-[MEM]-query content-addressable teacher pool #11

Merged
osobh merged 1 commits from d319-content-addressable-teacher into main 2026-06-17 16:32:58 +00:00
Conversation 0 Commits 1 Files Changed 1
+131 -8
1 changed files with 131 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/training/rtx-transformers/src/layers/set_encoder_teacher.rs
+130 -7
View File
@@ -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
// mean (order-invariant); recency mode uses an exponential-decay weight
// (recent tokens dominate) so the statistic tracks non-stationary state.
// Pool the L tokens → [1, d_model]. Learned-query mode lets a learned
// `[MEM]` query attend over the tokens (content-addressable: the encoder
// *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);