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 <noreply@anthropic.com>

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;

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<f32> = (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<Tensor> {
+    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<Tensor> {
+    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<RopeAttention> {
+        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<Tensor> {
+        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<Tensor> { 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<Layer>,
+    final_ln: LayerNorm,
+    tables: (Tensor, Tensor),
+}
+
+impl MemoryAttention {
+    pub fn load(w: &Weights) -> Result<MemoryAttention> {
+        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::<Result<Vec<_>>>()?;
+        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<Tensor> {
+        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::<f32>().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::<f32>().unwrap();
+        let s = sin.narrow(0, 0, 1).unwrap().to_vec2::<f32>().unwrap();
+        assert!(c[0].iter().all(|v| *v == 1.0) && s[0].iter().all(|v| *v == 0.0));
+    }
+}

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::<f32>().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}");
+}