ferritin_plms/esm3/layers/
relative_position.rs

1//! Relative position embeddings for the ESM3 structure encoder (ferritin-100.22).
2//!
3//! The structure encoder is *local*: each residue is encoded from its
4//! k-nearest neighbours in space, not from the whole chain. Those neighbours
5//! arrive in distance order, which throws away where they sat in the sequence
6//! — so the only thing telling the model that a spatial neighbour was also a
7//! sequence neighbour is this embedding.
8//!
9//! It is also the encoder's **initial hidden state**. There is no token
10//! embedding: the transformer starts from these relative offsets and lets all
11//! geometry enter through geometric attention.
12
13use candle_core::{DType, Result, Tensor};
14use candle_nn::{Embedding, Module, VarBuilder};
15
16/// Offsets are clamped to ±`BINS` before lookup.
17pub const BINS: i64 = 32;
18
19/// Table rows: `2 * BINS + 2`. The `+2` is one row per clamped extreme plus
20/// the padding index that shifts the range to be non-negative — which is why
21/// the checkpoint's table is 66 rows and not 65.
22pub const NUM_EMBEDDINGS: usize = (2 * BINS + 2) as usize;
23
24/// Embeds the signed sequence offset between a query residue and each of its
25/// spatial neighbours.
26#[derive(Debug)]
27pub struct RelativePositionEmbedding {
28    embedding: Embedding,
29}
30
31impl RelativePositionEmbedding {
32    /// Load the `(66, d_model)` table from `<prefix>.embedding.weight`.
33    pub fn load(vb: VarBuilder, d_model: usize) -> Result<Self> {
34        let weight = vb.get((NUM_EMBEDDINGS, d_model), "embedding.weight")?;
35        Ok(Self {
36            embedding: Embedding::new(weight, d_model),
37        })
38    }
39
40    /// `query`: `(N,)` centre residue indices. `key`: `(N, K)` neighbour
41    /// residue indices. Returns `(N, K, d_model)`.
42    ///
43    /// The offset is clamped to `±BINS` and then shifted by `BINS + 1`, so an
44    /// offset of `-BINS` lands on row 1 and row 0 stays reserved for padding.
45    pub fn forward(&self, query: &Tensor, key: &Tensor) -> Result<Tensor> {
46        let query = query.to_dtype(DType::I64)?.unsqueeze(1)?; // (N, 1)
47        let diff = key.to_dtype(DType::I64)?.broadcast_sub(&query)?;
48        let idx = diff.clamp(-BINS, BINS)?.affine(1.0, (BINS + 1) as f64)?;
49        self.embedding.forward(&idx.to_dtype(DType::U32)?)
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56    use candle_core::Device;
57
58    fn table(d: usize) -> Result<RelativePositionEmbedding> {
59        // Row i is filled with the value i, so a looked-up row reports its index.
60        let rows: Vec<f32> = (0..NUM_EMBEDDINGS)
61            .flat_map(|i| std::iter::repeat_n(i as f32, d))
62            .collect();
63        let w = Tensor::from_vec(rows, (NUM_EMBEDDINGS, d), &Device::Cpu)?;
64        Ok(RelativePositionEmbedding {
65            embedding: Embedding::new(w, d),
66        })
67    }
68
69    /// The checkpoint ships 66 rows; a 65-row table would mean the padding
70    /// offset was dropped and every lookup would be off by one.
71    #[test]
72    fn test_table_is_sixty_six_rows() {
73        assert_eq!(NUM_EMBEDDINGS, 66);
74    }
75
76    #[test]
77    fn test_zero_offset_lands_on_the_centre_row() -> Result<()> {
78        let t = table(1)?;
79        let q = Tensor::new(&[5i64], &Device::Cpu)?;
80        let k = Tensor::new(&[[5i64]], &Device::Cpu)?;
81        let out = t.forward(&q, &k)?.flatten_all()?.to_vec1::<f32>()?;
82        assert_eq!(out[0], (BINS + 1) as f32, "offset 0 -> row BINS+1");
83        Ok(())
84    }
85
86    /// Offsets beyond the window saturate rather than wrapping into a
87    /// different residue's row, and never reach the reserved padding row 0.
88    #[test]
89    fn test_offsets_clamp_and_never_hit_the_padding_row() -> Result<()> {
90        let t = table(1)?;
91        let q = Tensor::new(&[100i64], &Device::Cpu)?;
92        let k = Tensor::new(&[[0i64, 100 - 40, 100 + 40, 200]], &Device::Cpu)?;
93        let out = t.forward(&q, &k)?.flatten_all()?.to_vec1::<f32>()?;
94        assert_eq!(out[0], 1.0, "-100 clamps to -BINS -> row 1");
95        assert_eq!(out[1], 1.0, "-40 clamps to -BINS -> row 1");
96        assert_eq!(out[2], (2 * BINS + 1) as f32, "+40 clamps to +BINS");
97        assert_eq!(out[3], (2 * BINS + 1) as f32, "+100 clamps to +BINS");
98        for v in out {
99            assert_ne!(v, 0.0, "row 0 is reserved for padding");
100        }
101        Ok(())
102    }
103
104    #[test]
105    fn test_sign_is_preserved() -> Result<()> {
106        let t = table(1)?;
107        let q = Tensor::new(&[10i64], &Device::Cpu)?;
108        let k = Tensor::new(&[[9i64, 11]], &Device::Cpu)?;
109        let out = t.forward(&q, &k)?.flatten_all()?.to_vec1::<f32>()?;
110        assert_eq!(out[0], BINS as f32, "offset -1");
111        assert_eq!(out[1], (BINS + 2) as f32, "offset +1");
112        Ok(())
113    }
114}
ferritin_plms/esm3/layers/relative_position.rs

ferritin_plms/esm3/layers/
relative_position.rs
