Mechanisms of AI Protein Folding in ESMFold
This work provides interpretability insights for protein folding models, which is incremental for researchers in computational biology.
The study investigated how ESMFold folds proteins by analyzing its mechanisms on a beta hairpin, identifying two computational stages: early blocks initialize biochemical signals and late blocks develop spatial features, with interventions showing strong causal effects.
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.