From ea2899d0771be2c6f814a1d269a1d16f52e6151e Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:30:53 +0200 Subject: [PATCH] S2.22: memory attention with rotary encoding matches PyTorch Self and cross attention with axial 2-D RoPE (keys tiled over memory frames, object-pointer tokens left unrotated). Max abs error 2.6e-5. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/src/lib.rs | 1 + crates/sam2/src/memattn.rs | 173 +++++++++++++++++++++++++++++++++ crates/sam2/tests/reference.rs | 21 ++++ 3 files changed, 195 insertions(+) create mode 100644 crates/sam2/src/memattn.rs diff --git a/crates/sam2/src/lib.rs b/crates/sam2/src/lib.rs index 18fe5a8ae..53156c7ab 100644 --- a/crates/sam2/src/lib.rs +++ b/crates/sam2/src/lib.rs @@ -13,6 +13,7 @@ use std::fmt; pub mod decoder; pub mod hiera; +pub mod memattn; pub mod memory; pub mod model; pub mod neck; diff --git a/crates/sam2/src/memattn.rs b/crates/sam2/src/memattn.rs new file mode 100644 index 000000000..c3bfc0893 --- /dev/null +++ b/crates/sam2/src/memattn.rs @@ -0,0 +1,173 @@ +// Modified work: ported from Hugging Face `transformers` `models/sam2_video/modeling_sam2_video.py` +// (Sam2VideoMemoryAttention, Sam2VideoMemoryAttentionLayer, Sam2VideoRoPEAttention, +// Sam2VideoVisionRotaryEmbedding; Apache-2.0, Copyright The HuggingFace Inc. team) to Rust on +// candle. See NOTICE. +//! Memory attention: the current frame's features attend to themselves and to the memory bank +//! (memory frames plus object pointers) with axial 2-D rotary position encoding. + +use candle_core::{D, Device, Tensor}; +use candle_nn::{Linear, ops::softmax_last_dim}; + +use crate::nn::{LayerNorm, apply, layer_norm, linear}; +use crate::weights::Weights; +use crate::{Error, Result}; + +const DIM: usize = 256; +const LAYERS: usize = 4; +/// Side of the feature grid the rotary encoding is defined on. +pub const GRID: usize = 64; + +/// `cos` and `sin` of the axial rotary encoding for the `GRID x GRID` grid, each `[GRID², 256]`. +fn rope_tables(theta: f32) -> Result<(Tensor, Tensor)> { + let freqs = DIM / 2 / 2; // per axis: head_dim / 4 + let inv: Vec = (0..freqs).map(|i| 1.0 / theta.powf(i as f32 / freqs as f32)).collect(); + let n = GRID * GRID; + let (mut cos, mut sin) = (vec![0f32; n * DIM], vec![0f32; n * DIM]); + for p in 0..n { + let (x, y) = ((p % GRID) as f32, (p / GRID) as f32); + for j in 0..DIM / 2 { + let f = if j < freqs { x * inv.get(j).copied().unwrap_or(0.0) } else { y * inv.get(j - freqs).copied().unwrap_or(0.0) }; + for k in 0..2 { + if let (Some(c), Some(s)) = (cos.get_mut(p * DIM + 2 * j + k), sin.get_mut(p * DIM + 2 * j + k)) { + *c = f.cos(); + *s = f.sin(); + } + } + } + } + Ok((Tensor::from_vec(cos, (n, DIM), &Device::Cpu)?, Tensor::from_vec(sin, (n, DIM), &Device::Cpu)?)) +} + +/// `(x0, x1) -> (-x1, x0)` over adjacent pairs of the last dimension. +fn rotate_pairwise(x: &Tensor) -> Result { + let (n, d) = x.dims2()?; + let pairs = x.reshape((n, d / 2, 2))?; + let (x1, x2) = (pairs.narrow(2, 0, 1)?.squeeze(2)?, pairs.narrow(2, 1, 1)?.squeeze(2)?); + Ok(Tensor::stack(&[&x2.neg()?, &x1], D::Minus1)?.reshape((n, d))?) +} + +fn rope(x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result { + Ok(((x * cos)? + (rotate_pairwise(x)? * sin)?)?) +} + +struct RopeAttention { + q: Linear, + k: Linear, + v: Linear, + o: Linear, +} + +impl RopeAttention { + fn load(w: &Weights, prefix: &str, kv_dim: usize) -> Result { + Ok(RopeAttention { + q: linear(w, &format!("{prefix}.q_proj"), DIM, DIM)?, + k: linear(w, &format!("{prefix}.k_proj"), kv_dim, DIM)?, + v: linear(w, &format!("{prefix}.v_proj"), kv_dim, DIM)?, + o: linear(w, &format!("{prefix}.o_proj"), DIM, DIM)?, + }) + } + + /// `query`: `[N, 256]`, `key`/`value`: `[M, kv]`. The first `M - exclude` keys are rotated + /// (their table is `cos`/`sin` tiled to length); the last `exclude` (object pointers) are not. + fn forward(&self, query: &Tensor, key: &Tensor, value: &Tensor, tables: &(Tensor, Tensor), exclude: usize) -> Result { + let (cos, sin) = tables; + let q = rope(&apply(&self.q, query)?, cos, sin)?; + let k = apply(&self.k, key)?; + let v = apply(&self.v, value)?; + let m = k.dim(0)?; + let rotated = m.checked_sub(exclude).ok_or_else(|| Error::Input("more object pointers than memory tokens".into()))?; + let k = if rotated == 0 { + k + } else { + let n = cos.dim(0)?; + let reps = if rotated == n { + 1 + } else if rotated % n == 0 { + rotated / n + } else { + return Err(Error::Input(format!("memory length {rotated} is not a multiple of {n}"))); + }; + let tile = |t: &Tensor| -> Result { Ok(Tensor::cat(&vec![t; reps], 0)?) }; + let k_rot = rope(&k.narrow(0, 0, rotated)?, &tile(cos)?, &tile(sin)?)?; + if exclude == 0 { k_rot } else { Tensor::cat(&[&k_rot, &k.narrow(0, rotated, exclude)?], 0)? } + }; + let scale = (DIM as f64).powf(-0.5); + let att = softmax_last_dim(&((q * scale)?.matmul(&k.t()?.contiguous()?)?))?; + apply(&self.o, &att.matmul(&v)?) + } +} + +struct Layer { + self_attn: RopeAttention, + cross: RopeAttention, + ln1: LayerNorm, + ln2: LayerNorm, + ln3: LayerNorm, + lin1: Linear, + lin2: Linear, +} + +pub struct MemoryAttention { + layers: Vec, + final_ln: LayerNorm, + tables: (Tensor, Tensor), +} + +impl MemoryAttention { + pub fn load(w: &Weights) -> Result { + let layers = (0..LAYERS) + .map(|i| { + let p = format!("memory_attention.layers.{i}"); + Ok(Layer { + self_attn: RopeAttention::load(w, &format!("{p}.self_attn"), DIM)?, + cross: RopeAttention::load(w, &format!("{p}.cross_attn_image"), 64)?, + ln1: layer_norm(w, &format!("{p}.layer_norm1"), DIM, 1e-5)?, + ln2: layer_norm(w, &format!("{p}.layer_norm2"), DIM, 1e-5)?, + ln3: layer_norm(w, &format!("{p}.layer_norm3"), DIM, 1e-5)?, + lin1: linear(w, &format!("{p}.linear1"), DIM, 2048)?, + lin2: linear(w, &format!("{p}.linear2"), 2048, DIM)?, + }) + }) + .collect::>>()?; + Ok(MemoryAttention { layers, final_ln: layer_norm(w, "memory_attention.layer_norm", DIM, 1e-5)?, tables: rope_tables(10000.0)? }) + } + + /// Condition the current frame on memory. + /// + /// `current` and `current_pos`: `[4096, 256]` (the lowest-resolution neck level and its + /// position encoding, one row per cell); `memory` and `memory_pos`: `[M, 64]` (memory frames + /// followed by `num_pointers` object-pointer tokens). Returns `[4096, 256]`. + pub fn forward(&self, current: &Tensor, current_pos: &Tensor, memory: &Tensor, memory_pos: &Tensor, num_pointers: usize) -> Result { + if current.dims() != [GRID * GRID, DIM] { + return Err(Error::Input(format!("current features must be [{}, {DIM}], got {:?}", GRID * GRID, current.dims()))); + } + let keys_pos = (memory + memory_pos)?; + let mut out = (current + (current_pos * 0.1)?)?; + for l in &self.layers { + let q = l.ln1.forward(&out)?; + out = (&out + l.self_attn.forward(&q, &q, &q, &self.tables, 0)?)?; + let q = l.ln2.forward(&out)?; + out = (&out + l.cross.forward(&q, &keys_pos, memory, &self.tables, num_pointers)?)?; + let q = l.ln3.forward(&out)?; + out = (&out + apply(&l.lin2, &apply(&l.lin1, &q)?.relu()?)?)?; + } + self.final_ln.forward(&out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rotation_is_pairwise_and_tables_start_at_identity() { + let x = Tensor::new(&[[1f32, 2.0, 3.0, 4.0]], &Device::Cpu).unwrap(); + let r = rotate_pairwise(&x).unwrap().to_vec2::().unwrap(); + assert_eq!(r, vec![vec![-2.0, 1.0, -4.0, 3.0]]); + let (cos, sin) = rope_tables(10000.0).unwrap(); + // cell (0, 0): angle 0 everywhere + let c = cos.narrow(0, 0, 1).unwrap().to_vec2::().unwrap(); + let s = sin.narrow(0, 0, 1).unwrap().to_vec2::().unwrap(); + assert!(c[0].iter().all(|v| *v == 1.0) && s[0].iter().all(|v| *v == 0.0)); + } +} diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index f90d03ab8..d6d457871 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -173,3 +173,24 @@ fn memory_encoder_matches_pytorch() { eprintln!("memory position encoding: max abs error {e:e}"); assert!(e < 1e-5, "memory position error {e}"); } + +#[test] +fn memory_attention_matches_pytorch() { + let Some((w, r)) = load_video() else { return }; + let att = effectcraft_sam2::memattn::MemoryAttention::load(&w).unwrap(); + let get = |k: &str| r.get(k).unwrap().squeeze(1).unwrap(); + let pointers = r.get("mem.attn.num_pointers").unwrap().to_vec1::().unwrap()[0] as usize; + let out = att + .forward( + &get("mem.attn.current_vision_features"), + &get("mem.attn.current_vision_position_embeddings"), + &get("mem.attn.memory"), + &get("mem.attn.memory_posision_embeddings"), + pointers, + ) + .unwrap(); + let want = r.get("mem.attn.out").unwrap().squeeze(0).unwrap().squeeze(0).unwrap(); + let e = max_abs(&out, &want); + eprintln!("memory attention: max abs error {e:e} (pointers {pointers})"); + assert!(e < MEMORY_TOL, "memory attention error {e}"); +}