ferritin_plms/esm3/layers/
blocks.rs

1//! ESM3 UnifiedTransformerBlock: plain MHA + geometric attention + SwiGLU FFN.
2
3use crate::esm3::models::esm3::ESM3Config;
4use crate::esm3::utils::affine3d::Affine3D;
5use crate::esmc::layers::attention::MultiHeadAttention;
6use crate::esmc::layers::geom_attention::GeometricReasoningOriginalImpl;
7use crate::esmc::models::esmc::{ESMCConfig, ESMTokenizer, FfnType};
8use candle_core::{D, Module, Result, Tensor};
9use candle_nn::{self as nn, VarBuilder};
10
11// ── SwiGLU FFN ───────────────────────────────────────────────────────────────
12
13pub struct SwiGLU {
14    layer_norm: nn::LayerNorm,
15    linear1: nn::Linear,
16    linear2: nn::Linear,
17}
18
19impl SwiGLU {
20    fn hidden_dim(expansion_ratio: f64, d_model: usize) -> usize {
21        ((expansion_ratio * d_model as f64 + 255.0) / 256.0).floor() as usize * 256
22    }
23
24    /// Load the FFN, taking the bias terms **from the checkpoint** rather than
25    /// assuming they are absent (ferritin-100.27).
26    ///
27    /// The two checkpoints that share this block disagree, and hardcoding
28    /// either answer is silently wrong for the other:
29    ///
30    /// | checkpoint | `ffn.1.bias` / `ffn.3.bias` |
31    /// |---|---|
32    /// | `esm3_sm_open_v1.pth` (the sequence model) | absent |
33    /// | `esm3_structure_encoder_v0.pth` (the VQ-VAE encoder) | **present** |
34    ///
35    /// This used to be `linear_no_bias` unconditionally, which is correct for
36    /// the sequence model and drops two real tensors per block for the
37    /// structure encoder. Nothing caught it: the loader only ever asked for the
38    /// weights, so "every tensor resolves" stayed true while the biases sat
39    /// unread in the file. The cost was ~4% drift in the encoder's latent —
40    /// invisible to any self-consistency check, but enough to move 11 of 93
41    /// residues onto a different codebook entry.
42    ///
43    /// Probed rather than made a config flag because the checkpoint is the
44    /// authority here, and the same probing idiom already resolves ESM-C's
45    /// wrapper prefix in `loader::optional_prefix`.
46    pub fn load(vb: VarBuilder, config: &ESM3Config) -> Result<Self> {
47        let hidden = Self::hidden_dim(config.expansion_ratio, config.d_model);
48        let linear = |d_in, d_out, vb: VarBuilder| {
49            if vb.contains_tensor("bias") {
50                nn::linear(d_in, d_out, vb)
51            } else {
52                nn::linear_no_bias(d_in, d_out, vb)
53            }
54        };
55        Ok(Self {
56            layer_norm: nn::layer_norm(config.d_model, 1e-5, vb.pp("0"))?,
57            linear1: linear(config.d_model, hidden * 2, vb.pp("1"))?,
58            linear2: linear(hidden, config.d_model, vb.pp("3"))?,
59        })
60    }
61}
62
63impl Module for SwiGLU {
64    fn forward(&self, x: &Tensor) -> Result<Tensor> {
65        let x = self.layer_norm.forward(x)?;
66        let x = self.linear1.forward(&x)?;
67        let chunks = x.chunk(2, D::Minus1)?;
68        self.linear2.forward(&(chunks[0].silu()? * &chunks[1])?)
69    }
70}
71
72// ── ESM3 UnifiedTransformerBlock ─────────────────────────────────────────────
73
74pub struct UnifiedTransformerBlock {
75    /// Absent in the structure encoder, whose blocks are geometric-only.
76    attn: Option<MultiHeadAttention>,
77    geom_attn: Option<GeometricReasoningOriginalImpl>,
78    ffn: SwiGLU,
79    scaling_factor: f64,
80}
81
82impl UnifiedTransformerBlock {
83    pub fn load(vb: VarBuilder, config: &ESM3Config, layer_idx: usize) -> Result<Self> {
84        // Build an ESMCConfig shim so we can reuse the existing load() implementations.
85        let esmc_cfg = esmc_config_from_esm3(config);
86
87        let attn = Some(MultiHeadAttention::load(vb.pp("attn"), &esmc_cfg)?);
88
89        let geom_attn = if layer_idx < config.n_layers_geom {
90            // Checkpoints name this block "geom_attn", not "geometric"
91            // (ferritin-100.21).
92            Some(GeometricReasoningOriginalImpl::load(
93                vb.pp("geom_attn"),
94                &esmc_cfg,
95            )?)
96        } else {
97            None
98        };
99
100        let ffn = SwiGLU::load(vb.pp("ffn"), config)?;
101
102        Ok(Self {
103            attn,
104            geom_attn,
105            ffn,
106            scaling_factor: config.residue_scaling_factor(),
107        })
108    }
109
110    /// A geometric-only block, as used by the ESM3 structure encoder.
111    ///
112    /// The released `esm3_structure_encoder_v0.pth` has no `attn.*` tensors at
113    /// all — its two blocks are `geom_attn` plus `ffn`. Loading it through
114    /// [`load`][Self::load] fails on `attn.layernorm_qkv.0.weight`, which is
115    /// how this difference was found (ferritin-100.22).
116    pub fn load_geometric(vb: VarBuilder, config: &ESM3Config) -> Result<Self> {
117        let esmc_cfg = esmc_config_from_esm3(config);
118        let geom_attn = Some(GeometricReasoningOriginalImpl::load(
119            vb.pp("geom_attn"),
120            &esmc_cfg,
121        )?);
122        let ffn = SwiGLU::load(vb.pp("ffn"), config)?;
123        Ok(Self {
124            attn: None,
125            geom_attn,
126            ffn,
127            scaling_factor: config.residue_scaling_factor(),
128        })
129    }
130
131    pub fn forward(
132        &self,
133        x: &Tensor,
134        sequence_id: Option<&Tensor>,
135        affine: Option<&Affine3D>,
136        affine_mask: Option<&Tensor>,
137        chain_id: Option<&Tensor>,
138    ) -> Result<Tensor> {
139        let mut x = x.clone();
140
141        // Standard multi-head attention residual, when the block has one.
142        if let Some(attn) = &self.attn {
143            let r1 = attn.forward(&x, sequence_id)?;
144            x = (&x + (r1 / self.scaling_factor)?)?;
145        }
146
147        // Geometric attention residual (only in layers where geom_attn is present)
148        if let (Some(geom), Some(aff), Some(mask)) = (&self.geom_attn, affine, affine_mask) {
149            let r2 = geom.forward(&x, aff, mask, sequence_id, chain_id)?;
150            x = (&x + (r2 / self.scaling_factor)?)?;
151        }
152
153        // FFN residual
154        let r3 = self.ffn.forward(&x)?;
155        x = (&x + (r3 / self.scaling_factor)?)?;
156
157        Ok(x)
158    }
159}
160
161/// Build an `ESMCConfig` shim from `ESM3Config` so ESMC layer loaders can be reused.
162fn esmc_config_from_esm3(cfg: &ESM3Config) -> ESMCConfig {
163    let n_layers = cfg.n_layers;
164    ESMCConfig {
165        d_model: cfg.d_model,
166        n_heads: cfg.n_heads,
167        n_layers,
168        v_head_transformer: Some(cfg.v_head_transformer),
169        ffn_type: FfnType::SWIGLU,
170        tokenizer: ESMTokenizer::Esm3OpenSmall,
171        use_plain_attn: true,
172        n_layers_geom: cfg.n_layers_geom,
173        scale_residue: cfg.scale_residue,
174        residue_scaling_factor: cfg.residue_scaling_factor(),
175        mask_and_zero_frameless: cfg.mask_and_zero_frameless,
176        bias: cfg.bias,
177        qk_layernorm: cfg.qk_layernorm,
178        expansion_ratio: cfg.expansion_ratio,
179        // Unused by the layer loaders we call
180        regression_head_output_dim: 0,
181        regression_head_hidden_dim: 0,
182        embedding_dim: 0,
183    }
184}
185
186// ── Tests ─────────────────────────────────────────────────────────────────────
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191    use candle_core::{DType, Device};
192    use std::collections::HashMap;
193
194    fn cfg(d_model: usize) -> ESM3Config {
195        ESM3Config {
196            d_model,
197            expansion_ratio: 8.0 / 3.0,
198            ..ESM3Config::sm_open()
199        }
200    }
201
202    /// Build the tensors an FFN needs, optionally including the two Linear
203    /// biases the structure encoder ships and the sequence model does not.
204    fn ffn_tensors(d_model: usize, with_bias: bool, device: &Device) -> HashMap<String, Tensor> {
205        let hidden = SwiGLU::hidden_dim(8.0 / 3.0, d_model);
206        let mut ts = HashMap::new();
207        let ones = |dims: (usize, usize)| Tensor::ones(dims, DType::F32, device).unwrap();
208        ts.insert(
209            "0.weight".into(),
210            Tensor::ones(d_model, DType::F32, device).unwrap(),
211        );
212        ts.insert(
213            "0.bias".into(),
214            Tensor::zeros(d_model, DType::F32, device).unwrap(),
215        );
216        ts.insert("1.weight".into(), ones((hidden * 2, d_model)));
217        ts.insert("3.weight".into(), ones((d_model, hidden)));
218        if with_bias {
219            ts.insert(
220                "1.bias".into(),
221                Tensor::ones(hidden * 2, DType::F32, device).unwrap(),
222            );
223            ts.insert(
224                "3.bias".into(),
225                Tensor::ones(d_model, DType::F32, device).unwrap(),
226            );
227        }
228        ts
229    }
230
231    /// The two checkpoints sharing this block disagree about the FFN biases:
232    /// `esm3_sm_open_v1.pth` has none, `esm3_structure_encoder_v0.pth` has both.
233    /// Loading must follow the checkpoint, not a hardcoded choice
234    /// (ferritin-100.27).
235    ///
236    /// This is the regression guard for a bug that cost ~4% drift in the
237    /// structure encoder's latent and moved 11 of 93 residues onto a different
238    /// codebook entry — while every self-consistency check stayed green,
239    /// because the loader simply never asked for the tensors it was dropping.
240    #[test]
241    fn test_ffn_bias_follows_the_checkpoint() -> Result<()> {
242        let device = Device::Cpu;
243        let d_model = 16usize;
244
245        for with_bias in [false, true] {
246            let vb = VarBuilder::from_tensors(
247                ffn_tensors(d_model, with_bias, &device),
248                DType::F32,
249                &device,
250            );
251            let ffn = SwiGLU::load(vb, &cfg(d_model))?;
252            assert_eq!(
253                ffn.linear1.bias().is_some(),
254                with_bias,
255                "linear1 bias presence should follow the checkpoint (with_bias = {with_bias})"
256            );
257            assert_eq!(
258                ffn.linear2.bias().is_some(),
259                with_bias,
260                "linear2 bias presence should follow the checkpoint (with_bias = {with_bias})"
261            );
262        }
263        Ok(())
264    }
265
266    /// A checkpoint carrying biases must actually *use* them: the old
267    /// `linear_no_bias` loaded without error and silently ignored them, so
268    /// "it loads" was never evidence that it was right.
269    #[test]
270    fn test_ffn_biases_change_the_output() -> Result<()> {
271        let device = Device::Cpu;
272        let d_model = 16usize;
273        let x = Tensor::ones((1, 4, d_model), DType::F32, &device)?;
274
275        let plain = SwiGLU::load(
276            VarBuilder::from_tensors(ffn_tensors(d_model, false, &device), DType::F32, &device),
277            &cfg(d_model),
278        )?
279        .forward(&x)?;
280        let biased = SwiGLU::load(
281            VarBuilder::from_tensors(ffn_tensors(d_model, true, &device), DType::F32, &device),
282            &cfg(d_model),
283        )?
284        .forward(&x)?;
285
286        let a: Vec<f32> = plain.flatten_all()?.to_vec1()?;
287        let b: Vec<f32> = biased.flatten_all()?.to_vec1()?;
288        let diff = a
289            .iter()
290            .zip(&b)
291            .map(|(x, y)| (x - y).abs())
292            .fold(0f32, f32::max);
293        assert!(
294            diff > 1e-3,
295            "the biases should change the output; max diff {diff}"
296        );
297        Ok(())
298    }
299}
ferritin_plms/esm3/layers/blocks.rs

ferritin_plms/esm3/layers/
blocks.rs
