Structure, Not Sequence: Tracking and Validating Folding Mechanisms in ESMFold's Trunk with Sparse Autoencoders

Bohan Liu; Xiaohan Fan; Wenqian Ye; Aidong Zhang

q-bio.QMOct 26, 2026

Structure, Not Sequence: Tracking and Validating Folding Mechanisms in ESMFold's Trunk with Sparse Autoencoders

Authors: Bohan Liu, Xiaohan Fan, Wenqian Ye, Aidong Zhang

Organizations: University of Virginia Charlottesville VA 22903, USA

Abstract

Structure predictors like ESMFold fold proteins using neural network computation and have transformed structural biology. However, the driving mechanism from amino-acid (AA) sequence to 3D structure is insufficiently understood, and biological meaning remains locked in the dense network. Sparse autoencoders (SAEs) have shown promising results in obtaining interpretable concept features over an overcomplete basis of intermediate protein sequence-model representations. Yet folding models, which predict physical shape from AA sequence semantics, have no such account. We hypothesize that the features carrying the folding computation evolve in the 3D structural regression process. In this paper, we interpret the folding trunk in three steps. First, a pilot analysis on the representation distribution across modules and recycles locates distinct spatial computation in the trunk. Second, the dictionaries expose a division of labor: the single representation carries inherited semantics, and the pair representation computes the pairwise geometry. Third, we interpret the dictionary at scale with a reasoning language model, trace how the concepts compose across blocks by gradient-based attribution, and follow the interpreted code to the model's folding output, verifying where the computation aligns with, or departs from, known biochemistry. Together, the sparse dictionaries we train capture the relational geometry the trunk computes beyond the sequence model, and their concept activations forecast the model's own fold quality, turning the black-box folding trunk into a sparse, interpretable, and scientifically verifiable circuit.

Figures & tables

Panel of Figure 1 — Figure 1: Prior works interpret the ESM-2 sequence model trained on masked-language modeling, whereas we read ESMFold’s folding trunk: its single ( $s_{s}$ ) and pair ( $s_{z}$ ) representations, supervised by the 3D frame-aligned (FAPE) structure loss together with auxiliary read-out losses on the trunk states. The trunk and the structure head turn an amino-acid sequence embedding into a 3D structure with a per-residue confidence (pLDDT), with optional iterative refinement denoted as recycle in the figure.

Figure 2: Three pilot observations motivate the protocol. (A) Linear CKA between ESM-2 layers and trunk blocks. (B) Distance $R^{2}$ for pairwise residue geometry decoded from the pair representation $s_{z}$ versus any single-residue state $s_{s}$ , by sequence separation. (C) Detection auroc across trunk depth, from the frozen ESM-2 $L_{36}$ to $b_{47}$ , for a functional-site concept (catalytic) and a geometric one (burial).

Figure 3: The interpreted circuit among the annotation categories. The categories are the nodes, and integrated gradients route between them across the trunk’s depth. The pair representation

s_{z}

(top track) carries the geometry classes, the single representation

s_{s}

(bottom track) the residue-level categories. Two cross-track bundles carry the coupling on a shared width scale, with forward

s_{s}\!\to\!s_{z}

thick and backward

s_{z}\!\to\!s_{s}

thin.

Target	Type	ESM-2-3B	$+$ trunk	$\Delta$
Residue contact	relational	0.636	0.954	+0.318
Ladder continuation	relational	0.644	0.853	+0.209
Strand orientation	relational	0.873	0.959	+0.086
Metal binding	per-residue	0.978	0.992	+0.014
Strand	per-residue	0.961	0.969	+0.008
Active site	per-residue	0.978	0.984	+0.006

Table 1: What the trunk code adds to the language-model states prior dictionaries are trained on. Target : the property a probe is asked to predict. Type : whether the target is a relation between two residues or a property of one. ESM-2-3B : auroc of a probe on ESM-2-3B

L_{36}

alone, with the pairwise probe

[\,s_{i};\,s_{j};\,s_{i}\!\odot\!s_{j}\,]

for relational targets.

+

trunk : the same states concatenated with the trunk code.

\Delta

: the increment, to be read against a control that shuffles the trunk features across examples, which spans

\pm 0.010

on every row.

Figure 4: Activation sites of select concepts and proteins. The activation is clipped below 0.15 for clarity. Within-protein recall and precision are denoted as rec and prec respectively.

Edge	Source feature	Target feature	Explanation
supported by known biochemistry
$b_{12}\!\to\!b_{24}$ , fwd	Leucine	Long-range ${\sim}12$ Å C $\beta$ –C $\beta$	leucine tends to bury, and buried residues meet at contact distance
$b_{24}\!\to\!b_{36}$ , fwd	Exposed acidic loop residue	Long-range ${\sim}10$ Å proximity	acidic residues tend to be exposed, and exposed residues tend to pair
$b_{00}\!\to\!b_{12}$ , fwd	Active sites or termini	Long-range tertiary contact	a catalytic site requires residues far apart in sequence to meet
$b_{00}\!\to\!b_{12}$ , fwd	Glycine residue identity	Glycine-mediated contact	the target names glycine, the residue the source detects
rejected

Table 2: The eight chosen attribution cases. Edge : the attribution edge location, and whether it runs forward (

s_{s}\!\to\!s_{z}

) or backward (

s_{z}\!\to\!s_{s}

). Source and Target : the auto-interpreted features. Explanation : the biochemical explanation.

Figure 5: Interpreted features forecast fold quality (

n{=}70{,}859

residues,

452

proteins). (A) A ridge probe on the interpretable code predicts ESMFold’s per-residue confidence (pLDDT) and its true per-residue accuracy (lDDT against the reference structure) at every trunk block. (B) Structural concepts ranked by their association with true accuracy (Spearman

\rho

); green marks positive and red marks negative correlation.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Interpreted detectors painted on the folded structure across eight annotation types, four proteins each. Each panel colors a trunk detector’s per-residue map on the ESMFold structure with the annotated sites marked. Green labels mark well-localized features, and yellow labels mark distributed features.

Figure 7: The interpreted vocabulary of both representations. (A) Share of labeled features that the sequence determines and the share that geometry determines, for each representation under its own question and under the other’s. (B) Geometry share of the single-representation vocabulary at each trunk block. (C) Amino-acid purity by identity assignment, and top-context annotation density by functional-site assignment.

Figure 8: The same flow as Figure 3 resolved to individual concepts: the twelve largest single-representation and seven largest pair concepts, with the same two cross-track bundles, forward

s_{s}\!\to\!s_{z}

thick and backward

s_{z}\!\to\!s_{s}

thin. Node area is attributed mass, edge width summed weight, both log-scaled.

Block	$k{=}16$	$k{=}32$	$k{=}64$	$k{=}128$
b00	0.872	0.920	0.962	0.988
b12	0.900	0.934	0.963	0.984
b24	0.893	0.928	0.959	0.981
b36	0.863	0.908	0.945	0.973
b47	0.893	0.929	0.957	0.977

Appendix

Table 3: Reconstruction of the trunk single representation. Reported are explained variance on a held-out set. We use the sparsest setting,

k{=}16

, trading reconstruction for maximally interpretable codes.

		AA prior		raw auroc		within-AA auroc
Property	$n$	only	res.	naive	ctrl	naive	ctrl
Disulfide	771	0.995	C †	0.956	0.892	0.768	0.908
Glycosyl.	395	0.972	N ‡	0.950	0.691	0.563	0.782
Active site	1220	0.896	H	0.863	0.581	0.846	0.974
Metal bind.	926	0.894	D	0.792	0.605	0.759	0.924
Mod. res.	664	0.845	S	0.722	0.560	0.680	0.802

Appendix

Table 4: Amino-acid type acts both as a biological prior and as a bias on measurement. Rows are annotated site types. AA prior reports the auroc reachable from residue type alone together with the dominant residue, which the chemistry requires for a disulfide (

\dagger

) and the sequon for glycosylation (

\ddagger

Concept	Feat.	auroc	IoU	% sig.
Disulfide	5934	0.98	0.79	91
Carbohydrate	8061	0.95	0.13	52
Active site	2819	0.80	0.03	9
Signal peptide	6323	0.79	0.35	73
Transmembrane	5591	0.77	0.46	68
Metal binding	1314	0.76	0.48	76

Appendix

Table 5: Where each interpreted detector reads from. Rows are annotation types; columns give the feature the detection scan chose, its detection auroc , the mean overlap (IoU) between the feature’s top activating residues and the annotated ones, and the share of proteins where that overlap beats a residue-permutation baseline.

predictor	distance $R^{2}$
(representation regressed on distance)	$b_{47}$	$b_{24}$
$s_{z}$ dictionary code (pair SAE, TopK)	$0.69$	$0.64$
$s_{z}$ raw (pair vector, $128$ -d)	$0.67$	$0.62$
ESM-2-3B $L_{36}$ , pairwise probe	$-0.11$	$-0.11$
ESM-2-3B $L_{24}$ , pairwise probe	$-0.20$	$-0.20$
ESM-2-650M $L_{24}$ , pairwise probe	$-0.00$	$-0.00$

Appendix

Table 6: Out-of-sample distance

R^{2}

for predicting C

\alpha

–C

\alpha

distance from each representation (ridge, protein-stratified split,

60{,}000

pairs across

500

proteins). predictor : the representation supplying the features. distance

R^{2}

: coefficient of determination on held-out proteins; higher is better, negative when the fit underperforms predicting the mean. The two columns index the trunk block supplying the pair code (

b_{47}

b_{24}

); the single-residue baselines do not depend on the trunk block and repeat across columns. Pairwise probe denotes the interaction features

[\,s_{i};\,s_{j};\,s_{i}\!\odot\!s_{j}\,]

that give a per-residue representation a pairwise view of the pair

(i,j)

	$b_{24}$		$b_{47}$
Features labeled	$\beta$ -pair	helix $i,i{+}4$	$\beta$ -pair	helix $i,i{+}4$
$\beta$ -sheet / strand	0.938	0.839	0.932	0.900
$\alpha$ -helix	0.737	0.989	0.834	0.993
generic contact	0.829	0.943	0.803	0.901
random same-size set	0.675	0.665	0.683	0.742
full dictionary	0.965	0.992	0.960	0.993

Appendix

Table 7: Held-out auroc for predicting a geometric relation from the pair features carrying a given label. Features labeled : features grouped by the concept named in their auto-interpretation label, chosen blind to the target. Columns give the two relations, at two blocks.

Figure 9: The

s_{s}\!\to\!s_{s}

attribution, single-representation composition across depth, at the level of individual features. Nodes are features at a block, labeled by their own interpretation, sized by attributed mass and colored by category; edge width is the attributed weight. The title carries this direction’s share of the total attributed mass across the four.

Figure 10: The

s_{z}\!\to\!s_{s}

attribution, pair geometry read back into the single code, at the level of individual features. Nodes are features at a block, labeled by their own interpretation, sized by attributed mass and colored by category; edge width is the attributed weight. The title carries this direction’s share of the total attributed mass across the four.

Figure 11: The

s_{s}\!\to\!s_{z}

attribution, the single representation writing geometry, at the level of individual features. Nodes are features at a block, labeled by their own interpretation, sized by attributed mass and colored by category; edge width is the attributed weight. The title carries this direction’s share of the total attributed mass across the four.

Figure 12: The

s_{z}\!\to\!s_{z}

attribution, pair-geometry composition across depth, at the level of individual features. Nodes are features at a block, labeled by their own interpretation, sized by attributed mass and colored by category; edge width is the attributed weight. The title carries this direction’s share of the total attributed mass across the four.

Residue	geometry keywords credited to it
Gly	turn, local, bend, loop, flex, coil, hinge, termin, cap
Pro	turn, local, kink, bend, cap, loop, coil
Cys	disulf, cyste, bond, long-range, tertiary
Ile, Val	pack, buri, hydroph, sheet, strand, tertiary, long-range, core
Leu	pack, buri, hydroph, helix, tertiary, long-range, core
Ala	pack, buri, helix, hydroph, core

Appendix

Table 8: The rule set behind the grounding audit. An identity endpoint is credited when the opposite label contains one of its keywords; a structural endpoint is credited when the pair endpoint’s geometry class appears in its row.

Protein Structure Prediction

Protein structure prediction uses computational models to infer three-dimensional protein shapes and complexes from sequences and other biological data.

See all 40 papers

Two Stages of Folding: Convergent Mechanisms in AI Protein Folding Trunks
Feb 5, 2026Kevin Lu, Jannik Brinkmann, Stefan Huber +4Protein Structure Prediction AI-Assisted Scientific Research
1Northeastern University · 2TU Clausthal · 3Harvard University +2
Interpreting Latent Protein Language Model Features with Geometric Annotations
Aug 26, 2026Siddharth Setlur, Djordje Mihajlovic, Darrick LeeMechanistic Interpretability Protein Language Models
School of Mathematics University of Edinburgh, UK
PairSAE: Mechanistic Interpretability from Pair Representations in Protein Co-Folding
Jun 25, 2026Giosue Migliorini, Aristofanis Rontogiannis, Grigori Guitchounts +3Sparse Autoencoders Mechanistic Interpretability
University of California, Irvine · Flagship Pioneering

Feb 5, 2026J/K move · Enter open · S save

Two Stages of Folding: Convergent Mechanisms in AI Protein Folding Trunks

Kevin Lu, Jannik Brinkmann, Stefan Huber, Aaron Mueller+3

1Northeastern University · 2TU Clausthal · 3Harvard University+2

Protein Structure Prediction AI-Assisted Scientific Research Mechanistic Interpretability Representation Alignment

Open paper AskPDF arXiv

How do protein structure prediction models fold proteins? We investigate this question through causal interventions on the folding trunks of ESMFold, OpenFold, and Boltz-1. Across all three models, we find a shared two-stage computational structure. In the first stage, early blocks initialize pairwise biochemical signals: features like charge propagate from sequence into pairwise representations through architecture-specific pathways. In the second stage, late blocks develop pairwise spatial features: distance and contact information accumulate in the pairwise representation. We verify these mechanisms causally by showing that steering charge and distance features induces predictable structural changes. Furthermore, these representations are functionally interchangeable: pairwise states can be linearly aligned and substituted across models. Together, these results suggest that folding trunks with different architectures, inputs, and training procedures converge on a shared representational organization for mapping sequence chemistry into spatial geometry.

Abstract

Protein Structure Prediction

Explore similar work

Two Stages of Folding: Convergent Mechanisms in AI Protein Folding Trunks

Interpreting Latent Protein Language Model Features with Geometric Annotations

PairSAE: Mechanistic Interpretability from Pair Representations in Protein Co-Folding