ferritin_plms/esm3/layers/
output_heads.rs

1//! ESM3 output projection heads (OutputHeads).
2//!
3//! Projects the transformer hidden state to per-track logit distributions.
4
5use crate::esm3::models::esm3::ESM3Config;
6use candle_core::{Module, Result, Tensor};
7use candle_nn::{self as nn, LayerNormConfig, VarBuilder};
8
9// ── Generic regression head ───────────────────────────────────────────────────
10
11/// Linear(d_in) → GELU → LayerNorm → Linear(d_out).
12pub struct RegressionHead {
13    model: nn::Sequential,
14}
15
16impl RegressionHead {
17    pub fn load(vb: VarBuilder, d_in: usize, d_out: usize) -> Result<Self> {
18        let ln_conf = LayerNormConfig::from(1e-5);
19        let model = nn::seq()
20            .add(nn::linear(d_in, d_in, vb.pp("0"))?)
21            .add(nn::Activation::Gelu)
22            .add(nn::layer_norm(d_in, ln_conf, vb.pp("2"))?)
23            .add(nn::linear(d_in, d_out, vb.pp("3"))?);
24        Ok(Self { model })
25    }
26}
27
28impl Module for RegressionHead {
29    fn forward(&self, x: &Tensor) -> Result<Tensor> {
30        self.model.forward(x)
31    }
32}
33
34// ── OutputHeads ───────────────────────────────────────────────────────────────
35
36pub struct OutputHeads {
37    sequence_head: RegressionHead,
38    structure_head: RegressionHead,
39    ss8_head: RegressionHead,
40    sasa_head: RegressionHead,
41    function_head: RegressionHead,
42    residue_head: RegressionHead,
43    n_function_tracks: usize,
44    d_function_vocab: usize,
45}
46
47impl OutputHeads {
48    pub fn load(vb: VarBuilder, config: &ESM3Config) -> Result<Self> {
49        let d = config.d_model;
50        Ok(Self {
51            sequence_head: RegressionHead::load(
52                vb.pp("sequence_head"),
53                d,
54                config.d_sequence_vocab,
55            )?,
56            structure_head: RegressionHead::load(
57                vb.pp("structure_head"),
58                d,
59                config.d_structure_vocab,
60            )?,
61            ss8_head: RegressionHead::load(vb.pp("ss8_head"), d, config.d_ss8_vocab)?,
62            sasa_head: RegressionHead::load(vb.pp("sasa_head"), d, config.d_sasa_vocab)?,
63            function_head: RegressionHead::load(
64                vb.pp("function_head"),
65                d,
66                config.n_function_tracks * config.d_function_vocab,
67            )?,
68            residue_head: RegressionHead::load(vb.pp("residue_head"), d, config.d_residue_vocab)?,
69            n_function_tracks: config.n_function_tracks,
70            d_function_vocab: config.d_function_vocab,
71        })
72    }
73
74    /// Project hidden states to per-track logit distributions.
75    ///
76    /// `x`: `(B, L, d_model)` — post-norm transformer output.
77    ///
78    /// Returns `ESM3Output` with optional logit tensors.
79    pub fn forward(&self, x: &Tensor) -> Result<ESM3Output> {
80        let (b, l, _) = x.dims3()?;
81
82        // Function logits: (B, L, n_tracks * d_func_vocab) → (B, L, n_tracks, d_func_vocab)
83        let function_logits = self.function_head.forward(x)?.reshape((
84            b,
85            l,
86            self.n_function_tracks,
87            self.d_function_vocab,
88        ))?;
89
90        Ok(ESM3Output {
91            sequence_logits: Some(self.sequence_head.forward(x)?),
92            structure_logits: Some(self.structure_head.forward(x)?),
93            secondary_structure_logits: Some(self.ss8_head.forward(x)?),
94            sasa_logits: Some(self.sasa_head.forward(x)?),
95            function_logits: Some(function_logits),
96            residue_logits: Some(self.residue_head.forward(x)?),
97            embeddings: None,
98        })
99    }
100}
101
102// ── ESM3Output ────────────────────────────────────────────────────────────────
103
104/// Output of an ESM3 forward pass.
105#[derive(Debug)]
106pub struct ESM3Output {
107    /// `(B, L, d_sequence_vocab)` — per-residue amino-acid logits.
108    pub sequence_logits: Option<Tensor>,
109    /// `(B, L, d_structure_vocab)` — per-residue structure-token logits.
110    pub structure_logits: Option<Tensor>,
111    /// `(B, L, d_ss8_vocab)` — secondary-structure logits.
112    pub secondary_structure_logits: Option<Tensor>,
113    /// `(B, L, d_sasa_vocab)` — SASA-bin logits.
114    pub sasa_logits: Option<Tensor>,
115    /// `(B, L, n_function_tracks, d_function_vocab)` — function annotation logits.
116    pub function_logits: Option<Tensor>,
117    /// `(B, L, d_residue_vocab)` — InterPro residue-annotation logits.
118    pub residue_logits: Option<Tensor>,
119    /// `(B, L, d_model)` — final hidden states (when requested).
120    pub embeddings: Option<Tensor>,
121}
ferritin_plms/esm3/layers/output_heads.rs

ferritin_plms/esm3/layers/
output_heads.rs
