Mechanisms of AI Protein Folding in ESMFold

Abstract: How do protein structure prediction models fold proteins? We investigate this question by tracing how ESMFold folds a beta hairpin, a prevalent structural motif. Through counterfactual interventions on model latents, we identify two computational stages in the folding trunk. In the first stage, early blocks initialize pairwise biochemical signals: residue identities and associated biochemical features such as charge flow from sequence representations into pairwise representations. In the second stage, late blocks develop pairwise spatial features: distance and contact information accumulate in the pairwise representation. We demonstrate that the mechanisms underlying structural decisions of ESMFold can be localized, traced through interpretable representations, and manipulated with strong causal effects.

• The paper demonstrates that early ESMFold blocks propagate biochemical features via the seq2pair pathway, crucial for β-hairpin formation.
• Activation patching experiments identify two stages in the folding trunk: early biochemical feature integration and later geometric feature construction.
• The study establishes a causal link between latent charge representations and 3D structural outcomes, offering insights for rational protein design.

Mechanistic Dissection of ESMFold: Computational Stages in AI Protein Folding

Introduction

The emergence of deep learning-based models such as ESMFold has led to significant advances in single-sequence protein structure prediction. The paper "Mechanisms of AI Protein Folding in ESMFold" (2602.06020) undertakes a mechanistic interpretability analysis of ESMFold, focusing on how this architecture transitions from amino acid sequence to folded protein structure. Using targeted intervention techniques—primarily activation patching—the authors characterize a two-stage computational organization within the model's folding trunk and systematically localize the propagation of biochemical and geometric information across its layers. This interpretability-driven approach provides both empirical and causal insights into the model's operations.

ESMFold Architecture and Beta-Hairpin Motif

ESMFold utilizes a hierarchical architecture: a pretrained ESM-2 LLM encodes the input sequence, followed by a folding trunk that jointly refines a sequence representation and a pairwise representation across 48 iterative blocks, before a structure module decodes these to 3D atomic coordinates.

Figure 1: ESMFold’s architecture, illustrating the flow from protein LLM input to pairwise and sequence representations, through 48 folding trunk blocks, and finally to 3D structure output.

The analysis centers on the folding of the β-hairpin, a prevalent and well-characterized secondary structure motif comprising two antiparallel β-strands connected by a short loop.

Figure 2: 3D cartoon of a β-hairpin motif, highlighting the antiparallel strand arrangement and stabilizing backbone hydrogen bonds.

Activation Patching and Causal Localizations

The principal methodological advance is the application of “activation patching,” in which representations from a donor motif (β-hairpin) are transplanted into a corresponding region of a target protein during a single forward pass of the folding trunk. By manipulating which representation (sequence or pairwise) is patched and at what block depth, the study establishes the ability to induce or suppress specific secondary structural outcomes in the target.

Figure 3: Schematic of the activation patching setup—sequence and pairwise representations from a donor are integrated at specified blocks into the target’s folding trunk pathway.

Trunk patching of both representations across all blocks induces successful motif transfer in ~40% of cases, whereas input sequence or structure module interventions are markedly less effective. This finding robustly localizes folding decision computations to the folding trunk.

Two Computational Stages in the Folding Trunk

The key claim is the identification of two mechanistically distinct computational regimes:

1. Stage 1 (Early Blocks, 0–7): Biochemical Feature Propagation Early blocks propagate residue identity and associated biochemical features (notably charge) from the sequence representation (s) into the pairwise representation (z) via the seq2pair pathway. Sequence patching is exclusively effective in this window, as confirmed by successful induction rates and ablation studies.
2. Stage 2 (Late Blocks, 25+): Geometric Feature Construction The late blocks develop explicit spatial information in z (pairwise distances, contacts). Pairwise patching only becomes effective after block 25, reflecting that these layers support geometric rather than sequence-driven computations.

   Figure 4: Schematic of a folding block; clearly demarcates the Pair2Seq (affecting attention via $z$-derived biases) and Seq2Pair (injecting sequence-derived features into $z$) pathways.

Activation patching experiments, both forward (hairpin into helix) and reverse (helix into hairpin), reveal a symmetric two-stage pattern, supporting that this mechanism is a general property of the trunk architecture rather than a motif-specific artifact.

Linear Encoding and Causal Role of Biochemical Features

Charge Steering

Charge information is encoded in a prominent linear direction within the sequence representation. Difference-of-means analysis and linear probing confirm the linearly accessible separation of positively and negatively charged residues throughout trunk progression.

Manipulating the charge direction—by artificially steering two flanking helices in a target loop region to opposite charges—induces β-hairpin-like backbone hydrogen bonding, precisely during the early-stage window. This demonstrates a causal relationship between charge-like latent structure and the emergent fold.

Figure 5: Charge steering induces backbone H-bond formation; opposite-charge steering contracts the motif, same-charge steering increases cross-strand distance.

Propagation into Pairwise Representations

Linear probes show charge and residue identity information are transmitted via the seq2pair operator, becoming decodable from $z$ by block 10–15. Ablating this pathway abrogates successful motif transfer, corroborating the computational modularity of trunk information flow.

Emergence and Manipulation of Pairwise Geometric Features

Distance Decodability and Steering

By late blocks, linear regression probes achieve $R^2\approx 0.9$ on Cα–Cα distance prediction from $z$. Targeted interventions on $z$ using these probe weights (steering toward canonical cross-strand distances) produce direct modulation of 3D geometry.

Figure 6: Distance probes and steering—linear accessibility of geometric information in $z$ and selective hydrogen bond induction by steering at blocks 20–35.

Scaling the magnitude of pairwise embeddings at the structure module entry yields proportional expansions or contractions in the 3D output, confirming $z$ as a direct geometric blueprint.

Figure 7: Scaling the pairwise representation modulates global 3D structure size; $z$ functions as a distance matrix.

Pair2Seq Mechanism and Contact Promotion

The pair2seq pathway projects $z$ to a set of attn biases modulating sequence-level attention. In middle and late blocks, these bias values cleanly delineate structure contacts (ROC-AUC near 1.0), encouraging communication between residue pairs that are in close 3D proximity.

Figure 8: Pair2seq bias profile segregates structural contacts from non-contacts with high fidelity in late blocks.

Causal intervention studies—patching $z$—redistribute attention to favor donor contacts over target-specific contacts.

Figure 9: Changes in sequence attention probability post-pairwise patching, shifting focus toward donor contacts, confirming the causal influence of $z$ on attention.

Implications and Future Directions

The core implication is that ESMFold implements a modular computational stratification within the folding trunk: initial blocks as information integrators for biochemical features, later blocks as geometric solvers. The observed organization is robust to motif and directionality under patching experiments, indicative of a general design principle in deep learning-based structure predictors.

Contrary to a purely black-box or end-to-end learning paradigm, the causal evidence demonstrates that internal representations are linearly manipulable and admit targeted interventions with predictable, physically salient structural consequences. This demarcation between biochemical and geometric regimes enables controlled manipulation for interpretability, error diagnosis, and potentially, rational protein design.

Practically, these results advocate for the development of interpretable and steerable protein folding models, either through architectural modularity or the inclusion of explicit supervision for latent feature directions (e.g., charge, distance). The approach also offers prospects for integrating these mechanistic findings into higher-level tasks such as thermostability or functional site engineering, and for extending interpretability frameworks to other hierarchical architectures (AlphaFold2/Evoformer, protein-protein complexes).

Conclusion

The systematic intervention and probing analysis provided by "Mechanisms of AI Protein Folding in ESMFold" offers a high-resolution dissection of the model’s folding trunk. The work demonstrates that representational pathways in ESMFold are causally effective, interpretable, and robustly manipulable. The identification of two distinct computational regimes—early biochemical feature propagation and late geometric feature synthesis—provides actionable insight for mechanistic, diagnostic, and design-oriented extensions of protein structure prediction models.

(2602.06020)