Mechanisms of AI Protein Folding in ESMFold

TL;DR

This paper investigates how ESMFold predicts protein structures by tracing the model's internal mechanisms during folding, revealing two key computational stages and demonstrating interpretability and causal manipulation of its decisions.

Contribution

It identifies and characterizes the two main computational stages in ESMFold's folding process, providing insights into its internal mechanisms and interpretability.

Findings

Two computational stages in ESMFold's folding process

Residue identities and biochemical features are initialized early

Distance and contact information develop later

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.