ferritin_plms/esmfold2/
config.rs

1//! ESMFold2 configuration types.
2//!
3//! Defines [`ESMFold2Config`], which captures every hyper-parameter of the
4//! ESMFold2-Fast architecture. All defaults come from the official
5//! `biohub/ESMFold2-Fast` `config.json` and the architecture audit (2026-05-27).
6
7/// Full configuration for ESMFold2-Fast (biohub/ESMFold2-Fast).
8///
9/// All values default to the ESMFold2-Fast configuration unless noted.
10/// Source: biohub/ESMFold2-Fast config.json + architecture audit (2026-05-27).
11#[derive(Debug, Clone)]
12pub struct ESMFold2Config {
13    // ---- ESMC-6B backbone ----
14    /// HuggingFace repo for the frozen ESMC-6B encoder backbone.
15    pub esmc_model_id: &'static str, // "biohub/ESMC-6B"
16    pub lm_d_model: usize,    // 2560
17    pub lm_num_layers: usize, // 80
18
19    // ---- LM adapter (per recycling loop) ----
20    pub lm_encoder_n_layers: usize, // 4  (projects 2560 → d_single=384)
21    pub lm_dropout: f64,            // 0.25
22
23    // ---- Token representation ----
24    pub d_inputs: usize, // 451  (raw input feature dim)
25    pub d_single: usize, // 384  (single/sequence repr dim after LM adapter)
26    pub d_pair: usize,   // 256  (pair repr dim)
27
28    // ---- Relative position encoding ----
29    pub n_relative_residx_bins: usize, // 32
30    pub n_relative_chain_bins: usize,  // 2
31
32    // ---- Atom encoder ----
33    pub d_atom: usize,                  // 128
34    pub d_token_atom: usize,            // 768  (atom encoder output, projected to d_token)
35    pub atom_encoder_n_blocks: usize,   // 3
36    pub atom_encoder_n_heads: usize,    // 4
37    pub atom_encoder_swa_window: usize, // 128
38
39    // ---- Folding trunk ----
40    pub trunk_n_layers: usize, // 24
41    pub trunk_n_heads: usize,  // 8
42    pub trunk_dropout: f64,    // 0.25
43
44    // ---- MSA encoder (disabled in Fast) ----
45    pub msa_enabled: bool, // false for ESMFold2-Fast
46
47    // ---- Diffusion module ----
48    pub c_token: usize,          // 768   (token dim inside diffusion)
49    pub token_num_blocks: usize, // 12
50    pub token_num_heads: usize,  // 16
51    pub c_atom: usize,           // 128
52    pub atom_num_blocks: usize,  // 3
53    pub atom_num_heads: usize,   // 4
54    pub fourier_dim: usize,      // 256   (noise level encoding)
55    pub sigma_data: f64,         // 16.0  Å
56    // Inference noise schedule
57    pub inference_s_max: f64,       // 160.0
58    pub inference_s_min: f64,       // 0.0004
59    pub inference_num_steps: usize, // 14   (default; use 50 for quality)
60    pub inference_p: f64,           // 7.0
61    pub noise_scale: f64,           // 1.003
62    pub gamma_0: f64,               // 0.8
63    pub gamma_min: f64,             // 1.0
64    pub step_scale: f64,            // 1.5
65
66    // ---- Confidence head ----
67    pub confidence_n_layers: usize, // 4
68    pub num_plddt_bins: usize,      // 50
69    pub num_pae_bins: usize,        // 64
70    pub num_pde_bins: usize,        // 64
71    pub distogram_bins: usize,      // 39
72}
73
74impl ESMFold2Config {
75    /// Returns the ESMFold2-Fast configuration.
76    ///
77    /// Values are taken from `biohub/ESMFold2-Fast` `config.json` and the
78    /// architecture audit dated 2026-05-27. Use `inference_num_steps = 50`
79    /// (instead of the default 14) when higher structure quality is desired.
80    pub fn fast() -> Self {
81        Self {
82            // ESMC-6B backbone
83            esmc_model_id: "biohub/ESMC-6B",
84            lm_d_model: 2560,
85            lm_num_layers: 80,
86
87            // LM adapter
88            lm_encoder_n_layers: 4,
89            lm_dropout: 0.25,
90
91            // Token representation
92            d_inputs: 451,
93            d_single: 384,
94            d_pair: 256,
95
96            // Relative position encoding
97            n_relative_residx_bins: 32,
98            n_relative_chain_bins: 2,
99
100            // Atom encoder
101            d_atom: 128,
102            d_token_atom: 768,
103            atom_encoder_n_blocks: 3,
104            atom_encoder_n_heads: 4,
105            atom_encoder_swa_window: 128,
106
107            // Folding trunk
108            trunk_n_layers: 24,
109            trunk_n_heads: 8,
110            trunk_dropout: 0.25,
111
112            // MSA encoder disabled in Fast variant
113            msa_enabled: false,
114
115            // Diffusion module
116            c_token: 768,
117            token_num_blocks: 12,
118            token_num_heads: 16,
119            c_atom: 128,
120            atom_num_blocks: 3,
121            atom_num_heads: 4,
122            fourier_dim: 256,
123            sigma_data: 16.0,
124
125            // Inference noise schedule
126            inference_s_max: 160.0,
127            inference_s_min: 0.0004,
128            inference_num_steps: 14,
129            inference_p: 7.0,
130            noise_scale: 1.003,
131            gamma_0: 0.8,
132            gamma_min: 1.0,
133            step_scale: 1.5,
134
135            // Confidence head
136            confidence_n_layers: 4,
137            num_plddt_bins: 50,
138            num_pae_bins: 64,
139            num_pde_bins: 64,
140            distogram_bins: 39,
141        }
142    }
143}
ferritin_plms/esmfold2/config.rs

ferritin_plms/esmfold2/
config.rs
