Mechanisms of AI Protein Folding in ESMFold

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.

💡 Research Summary

This paper tackles a fundamental question in modern protein‑structure prediction: how does a neural network actually “fold” a protein? The authors focus on ESMFold, a state‑of‑the‑art model that, like AlphaFold2, uses a large language model to embed the amino‑acid sequence and then refines both a per‑residue (sequence) representation and a pairwise (residue‑by‑residue) representation through a 48‑layer folding trunk before decoding coordinates with a structure module.

To probe the internal computation, the authors adopt activation‑patching, a counterfactual intervention technique originally used in language‑model interpretability. They select a donor protein that contains a β‑hairpin (two short antiparallel strands linked by a turn) and a target protein that is predominantly α‑helical. For each of the 48 trunk blocks they extract the donor’s sequence vectors s and pairwise tensors z for the hairpin region, then replace the corresponding vectors in the target’s forward pass with either the donor’s s (sequence patch) or z (pairwise patch). After the patched forward pass they ask whether the resulting 3‑D structure contains a hairpin at the intervened region, using DSSP secondary‑structure assignment as the metric.

Across roughly 5 000 patching experiments (≈95 targets × 10 donors per loop × multiple loops), full‑patching (both s and z at every block) succeeds about 40 % of the time, confirming that the trunk is the most effective locus for structural manipulation. When the intervention is restricted to a single block, a striking two‑phase pattern emerges. Sequence patches are highly effective in the earliest blocks (0‑7), with success rates peaking near 40 % at block 0 and then rapidly declining. Pairwise patches, by contrast, become increasingly effective only after block ≈ 25, reaching about 20 % success by block 35. The same pattern appears when the direction of the experiment is reversed (patching helical representations into a hairpin target), indicating that the phenomenon reflects a general processing pipeline rather than a hairpin‑specific quirk.

The authors then dissect the two phases mechanistically. Early blocks are dominated by the seq2pair pathway: the sequence representation is projected into a pairwise update via element‑wise multiplication and subtraction (Eq. 3). This operation injects biochemical attributes—charge, polarity, side‑chain volume—into the pairwise tensor, effectively building a “chemical interaction map” before any explicit spatial information is present. Empirically, after a sequence patch at block 0 the patched z aligns closely with the donor’s z (as measured by an interpolation coefficient α close to 1). Freezing the seq2pair operation abolishes this alignment and dramatically reduces hairpin formation, confirming its causal role.

Later blocks are dominated by the pairwise‑to‑sequence (pair2seq) pathway and the triangular updates that follow. By this stage the pairwise tensor already encodes rich biochemical cues; the model now uses multiplicative attention and triangular attention to convert these cues into distance‑like signals. The resulting pairwise values are projected back into scalar biases for each attention head (Eq. 2), modulating the sequence self‑attention so that residues that should be close in 3‑D space attend more strongly to each other. This “spatial feature construction” directly shapes the coordinates produced by the downstream structure module. Pairwise patches applied in late blocks succeed because they overwrite the emergent distance map, steering the model toward the desired hairpin geometry.

Four concrete contributions are claimed: (1) establishing activation‑patching as a practical causal tool for protein‑folding models; (2) identifying two distinct computational stages—early biochemical signal propagation and later spatial feature construction—and assigning interpretable functions to each; (3) demonstrating that molecular properties such as charge are encoded in identifiable directions in latent space and can be steered to induce structural changes; (4) showing that the pairwise representation functions as a distance map that directly controls output geometry.

The paper concludes that ESMFold’s internal representations are not opaque black boxes but admit localized, causal, and interpretable interventions. This opens avenues for debugging folding failures, designing more modular architectures (e.g., separating chemistry‑only and geometry‑only modules), and potentially guiding de‑novo protein design by directly manipulating latent biochemical or geometric cues. Future work is suggested to extend the methodology to larger motifs (β‑α‑β loops, domain interfaces) and to explore whether the identified two‑stage pipeline holds across other transformer‑based folding models.