ESMFold Architecture: Protein Folding Model

• ESMFold is a deep learning protein structure prediction model that converts amino acid sequences into detailed 3D architectures without using MSAs or templates.
• The model employs an ESM-2 encoder, a 48-block folding trunk with dual sequence and pairwise updates, and an IPA decoder to output atomic coordinates.
• Causal tracing and counterfactual interventions enable precise mapping of biochemical features to spatial geometry, enhancing interpretability and control.

ESMFold is a deep learning-based protein structure prediction architecture built around a compositional “sequence → pair → structure” paradigm. Distinct from models reliant on input MSAs or structural templates, ESMFold leverages large-scale protein language modeling for its sequence encoding, and employs a trunk with explicit two-phase computation to successively transduce biochemical information into spatially precise atomic coordinates. Its structure enables causal, interpretable interventions at various representational levels, making it a mechanism-oriented system for protein folding research and application (Lu et al., 5 Feb 2026).

1. High-Level Pipeline and Architectural Organization

The ESMFold pipeline operates in three discrete stages:

1. ESM-2 Encoder: The raw amino-acid sequence $x = (x_1,\dots,x_L)$ is mapped to per-residue representations $s^{(0)} \in \mathbb{R}^{L\times d_s}$ using a pretrained protein LLM. ESM-2 employs 33 transformer layers, $d_s \simeq 1280$, $H_e=20$ heads, and an MLP width of about $4d_s$.
2. Folding Trunk: A sequence of $K=48$ identical blocks iteratively refines two intertwined latent states: $s^{(k)} \in \mathbb{R}^{L\times d_s}$ (sequence embedding) and $z^{(k)} \in \mathbb{R}^{L\times L\times d_z}$ (pairwise “distance map”). Each block encodes two sublayers: a sequence update with pairwise-attention bias, and a pairwise update with sequence-diffusion and triangular operators.
3. Structure Module: An Invariant Point Attention (IPA) decoder runs for 8 passes, consuming $(s^{(48)}, z^{(48)})$ and outputting $X \in \mathbb{R}^{L\times 14\times 3}$, i.e., coordinates for up to 14 atoms per residue.

ESMFold can also “recycle” trunk blocks across multiple passes (up to four), although mechanistic studies typically disable recycling for intervention clarity (Lu et al., 5 Feb 2026).

2. Detailed Structure of the Folding Trunk

Each trunk block $g^{(k)}$ contains two principal computational stages linked by residual connections:

• Sequence Update ($g_1$): This sublayer incorporates pair2seq bias, where each self-attention head $h$ at indices $(i, j)$ receives an additive term $\beta_{ijh}(z_{ij}) = [W_\beta z_{ij}]_h$. This term allows the current pairwise geometry $z_{ij}$ to directly modulate the per-residue attention patterns, implemented via multihead attention and a subsequent MLP.
• Pairwise Update ($g_2$): This sublayer performs seq2pair injection. Sequence information is injected into pairwise representation by projecting $s_i$ and $s_j$ into $u_i=W_u s_i$ and $v_j=W_v s_j$, then forming $$\varphi_{ij} = [u_i \odot v_j; u_i - v_j] \in \mathbb{R}^{2d_z}$$ and updating $z_{ij}$ via $z'_{ij} = z_{ij} + W_z \varphi_{ij}$. This is followed by triangular multiplicative updates and triangular self-attention, sharing design principles with AlphaFold’s Evoformer.

The block-level architecture remains constant, but the model’s representational regimes shift markedly across trunk depth, as revealed by causal activation analyses (Lu et al., 5 Feb 2026).

3. Two-Stage Representational Regimes: Biochemistry and Geometry

Empirical and intervention-based analyses reveal two computational stages in the folding trunk:

• Early "Biochemistry-Write" Stage ($k=0$ to $7$):
  • Dominated by sequence-to-pair (seq2pair) transfer.
  • High relative contribution from seq2pair updates; pair2seq influence is weak.
  • Biochemical features—including amino-acid identity and charge—are copied from sequence space into pairwise space. For instance, a direction $v_{\text{charge}}$ can be linearly separated in $s$ by averaging embeddings over positively (K, R, H) and negatively charged (D, E) residues. Perturbing $s$ along $v_{\text{charge}}$ alters $z$ via seq2pair, steering structural motifs such as electrostatic complementarity.
• Late "Geometry-Refinement" Stage ($k=25$ to $47$):
  • Pairwise-to-sequence (pair2seq) influence strengthens.
  • Triangular operators on $z$ crystallize a faithful “distance map” encoding, enabling probe-based and gradient-based manipulation of prospective spatial contacts. The Cα–Cα distance can be linearly decoded from $z_{ij}$ with $R^2$ rising from $\approx 0.2$ at block $0$ to $\approx 0.9$ at block $48$.
  • Contact information in $z_{ij}$, projected via $W_\beta$ in pair2seq, yields a per-head attention bias $\beta_{ijh}$ that segregates spatial contacts (Cα–Cα $<$ 8 Å) from non-contacts (ROC-AUC ≈ 0.95).
  • Pair2seq in late blocks focuses sequence attention on spatially proximal residues, reinforcing formation of correct tertiary contacts.

4. Causal Tracing and Counterfactual Intervention

To mechanistically localize computation, ESMFold supports activation patching—direct latent substitution at arbitrary block depth:

• Let $\ell_k$ denote the latent ($s^{(k)}$ or $z^{(k)}$) at block $k$ for a target.
• Substituting a region from a donor at layer $k$, one can propagate forward to $k=48$ and quantify:
  • Structural success (DSSP-based motif formation rate)
  • Representation alignment $$\alpha = \langle z_{\text{patched}} - z_{\text{target}},\, z_{\text{donor}} - z_{\text{target}}\rangle / \|z_{\text{donor}} - z_{\text{target}}\|^2$$

Findings:

• Sequence patching is effective only for $k\in[0,7]$ (40% success at $k=0$), matching the biochemistry write-in window.
• Pairwise patching impacts only $k\geq 25$ (peak ≈ 20% at $k \approx 35$).
• Freezing seq2pair updates in blocks $0$–$10$ eliminates sequence patching effects, confirming that seq2pair is the channel for early biochemical transfer.
• Ablating pair2seq in late blocks attenuates, but does not abolish, the effect of pairwise patching, suggesting additional, possibly parallel, geometric information pathways.

5. Architectural Hyperparameters

The following summarizes ESMFold’s principal numerical and representational settings:

Component	Layers/Blocks	Dimensions	Notable Settings
ESM-2 Encoder	33	$d_s \simeq 1280$	$H_e = 20$ heads, MLP width $\simeq 4d_s$
Folding Trunk	48	$d_z \simeq 256$	$H = 8$ attention heads, MLP width $\simeq 4d_s$
Structure Module	8 IPA passes	$L\times 14\times 3$	Up to 14 atoms/residue, no MSAs/templates input

ESMFold does not employ MSA-based featurization or template encodings at any stage. The initial pairwise $z^{(0)}$ is set via a learned relative positional embedding of residue index differences ($i-j$).

6. Formal Outline and Pseudocode

The high-level procedure of ESMFold’s folding trunk can be schematized as:

s⁰ = ESM2_Encoder(x) # ℝ^{L×d_s} z⁰ = PosEmbedding(i–j) # ℝ^{L×L×d_z} for k in 1…48: # Sequence Update (pair2seq + self-attention + MLP) for each head h, positions i,j: bias_{ijh} = [W_β · z^{k–1}_{ij}]_h A_{ij}^{(h)} = ⟨q_i, k_j⟩/√d_h + bias_{ijh} s̃ = MultiHeadAttention(s^{k–1}, scores=A) ŝ = s^{k–1} + s̃ s^k = ŝ + MLP(ŝ) # Pairwise Update (seq2pair + triangular updates + MLP) for all i,j: u_i = W_u · s_i^k v_j = W_v · s_j^k φ_{ij} = concat(u_i⊙v_j, u_i–v_j) z'_{ij} = z^{k–1}_{ij} + W_z · φ_{ij} z″ = TriangularMultiplicative(z') z̄ = TriangularSelfAttention(z″) z^k = z' + z̄ + MLP(z̄) X = IPA(s^{48}, z^{48}) return X # predicted atom coordinates

Blocks $0$–$7$ correspond to the biochemistry-writing phase; blocks $25$–$47$ encode geometric refinement.

7. Significance and Distinctive Mechanistic Features

ESMFold substantiates that protein folding can be approached as a sequence of write-in and refinement phases, where explicit modes of representation transfer (seq2pair and pair2seq) act as mechanistically localizable bridges between biochemical features and three-dimensional geometry. Activation patching and gradient-based steering demonstrate that these phases not only correlate with folding outcomes but have robust, causal effects on structure formation. The architecture accommodates transparent tracing and manipulation, thereby offering an empirical testbed for mechanistic protein folding hypotheses beyond the “black-box” regime (Lu et al., 5 Feb 2026).

A plausible implication is that such mechanism-localizing designs, leveraging explicit representational pipelines and causal interventional tools, may serve as templates for future protein structure predictors and for interpretability research in biological AI systems.