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    pub fn load(vb: VarBuilder, config: &ESM3Config) -> Result<Self> {
25        let hidden = Self::hidden_dim(config.expansion_ratio, config.d_model);
26        Ok(Self {
27            layer_norm: nn::layer_norm(config.d_model, 1e-5, vb.pp("0"))?,
28            linear1: nn::linear_no_bias(config.d_model, hidden * 2, vb.pp("1"))?,
29            linear2: nn::linear_no_bias(hidden, config.d_model, vb.pp("3"))?,
30        })
31    }
32}
33
34impl Module for SwiGLU {
35    fn forward(&self, x: &Tensor) -> Result<Tensor> {
36        let x = self.layer_norm.forward(x)?;
37        let x = self.linear1.forward(&x)?;
38        let chunks = x.chunk(2, D::Minus1)?;
39        self.linear2.forward(&(chunks[0].silu()? * &chunks[1])?)
40    }
41}
42
43// ── ESM3 UnifiedTransformerBlock ─────────────────────────────────────────────
44
45pub struct UnifiedTransformerBlock {
46    attn: MultiHeadAttention,
47    geom_attn: Option<GeometricReasoningOriginalImpl>,
48    ffn: SwiGLU,
49    scaling_factor: f64,
50}
51
52impl UnifiedTransformerBlock {
53    pub fn load(vb: VarBuilder, config: &ESM3Config, layer_idx: usize) -> Result<Self> {
54        // Build an ESMCConfig shim so we can reuse the existing load() implementations.
55        let esmc_cfg = esmc_config_from_esm3(config);
56
57        let attn = MultiHeadAttention::load(vb.pp("attn"), &esmc_cfg)?;
58
59        let geom_attn = if layer_idx < config.n_layers_geom {
60            Some(GeometricReasoningOriginalImpl::load(
61                vb.pp("geometric"),
62                &esmc_cfg,
63            )?)
64        } else {
65            None
66        };
67
68        let ffn = SwiGLU::load(vb.pp("ffn"), config)?;
69
70        Ok(Self {
71            attn,
72            geom_attn,
73            ffn,
74            scaling_factor: config.residue_scaling_factor(),
75        })
76    }
77
78    pub fn forward(
79        &self,
80        x: &Tensor,
81        sequence_id: Option<&Tensor>,
82        affine: Option<&Affine3D>,
83        affine_mask: Option<&Tensor>,
84        chain_id: Option<&Tensor>,
85    ) -> Result<Tensor> {
86        let mut x = x.clone();
87
88        // Standard multi-head attention residual
89        let r1 = self.attn.forward(&x, sequence_id)?;
90        x = (&x + (r1 / self.scaling_factor)?)?;
91
92        // Geometric attention residual (only in layers where geom_attn is present)
93        if let (Some(geom), Some(aff), Some(mask)) = (&self.geom_attn, affine, affine_mask) {
94            let r2 = geom.forward(&x, aff, mask, sequence_id, chain_id)?;
95            x = (&x + (r2 / self.scaling_factor)?)?;
96        }
97
98        // FFN residual
99        let r3 = self.ffn.forward(&x)?;
100        x = (&x + (r3 / self.scaling_factor)?)?;
101
102        Ok(x)
103    }
104}
105
106/// Build an `ESMCConfig` shim from `ESM3Config` so ESMC layer loaders can be reused.
107fn esmc_config_from_esm3(cfg: &ESM3Config) -> ESMCConfig {
108    let n_layers = cfg.n_layers;
109    ESMCConfig {
110        d_model: cfg.d_model,
111        n_heads: cfg.n_heads,
112        n_layers,
113        v_head_transformer: Some(cfg.v_head_transformer),
114        ffn_type: FfnType::SWIGLU,
115        tokenizer: ESMTokenizer::Esm3OpenSmall,
116        use_plain_attn: true,
117        n_layers_geom: cfg.n_layers_geom,
118        scale_residue: cfg.scale_residue,
119        residue_scaling_factor: cfg.residue_scaling_factor(),
120        mask_and_zero_frameless: cfg.mask_and_zero_frameless,
121        bias: cfg.bias,
122        qk_layernorm: cfg.qk_layernorm,
123        expansion_ratio: cfg.expansion_ratio,
124        // Unused by the layer loaders we call
125        regression_head_output_dim: 0,
126        regression_head_hidden_dim: 0,
127        embedding_dim: 0,
128    }
129}
ferritin_plms/esm3/layers/blocks.rs

ferritin_plms/esm3/layers/
blocks.rs
