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
Figure 1: Prior works interpret the ESM-2 sequence model trained on masked-language modeling, whereas we read ESMFold’s folding trunk: its single ( ss ) and pair ( sz ) 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 R2 for pairwise residue geometry decoded from the pair representation sz versus any single-residue state ss , by sequence separation. (C) Detection auroc across trunk depth, from the frozen ESM-2 L36 to b47 , 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 sz (top track) carries the geometry classes, the single representation ss (bottom track) the residue-level categories. Two cross-track bundles carry the coupling on a shared width scale, with forward ss→sz thick and backward sz→ss thin.
Target
Type
ESM-2-3B
+ trunk
Δ
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 L36 alone, with the pairwise probe [si;sj;si⊙sj] for relational targets. + trunk : the same states concatenated with the trunk code. Δ : the increment, to be read against a control that shuffles the trunk features across examples, which spans ±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
b12→b24 , fwd
Leucine
Long-range ∼12 Å C β –C β
leucine tends to bury, and buried residues meet at contact distance
b24→b36 , fwd
Exposed acidic loop residue
Long-range ∼10 Å proximity
acidic residues tend to be exposed, and exposed residues tend to pair
b00→b12 , fwd
Active sites or termini
Long-range tertiary contact
a catalytic site requires residues far apart in sequence to meet
b00→b12 , 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 ( ss→sz ) or backward ( sz→ss ). 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 ρ ); 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 ss→sz thick and backward sz→ss 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 ( † ) and the sequon for glycosylation ( ‡ ).
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 R2
(representation regressed on distance)
b47
b24
sz dictionary code (pair SAE, TopK)
0.69
0.64
sz raw (pair vector, 128 -d)
0.67
0.62
ESM-2-3B L36 , pairwise probe
−0.11
−0.11
ESM-2-3B L24 , pairwise probe
−0.20
−0.20
ESM-2-650M L24 , pairwise probe
−0.00
−0.00
Appendix
Table 6: Out-of-sample distance R2 for predicting C α –C α distance from each representation (ridge, protein-stratified split, 60,000 pairs across 500 proteins). predictor : the representation supplying the features. distance R2 : 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 ( b47 , b24 ); the single-residue baselines do not depend on the trunk block and repeat across columns. Pairwise probe denotes the interaction features [si;sj;si⊙sj] that give a per-residue representation a pairwise view of the pair (i,j) .
b24
b47
Features labeled
β -pair
helix i,i+4
β -pair
helix i,i+4
β -sheet / strand
0.938
0.839
0.932
0.900
α -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 ss→ss 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 sz→ss 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 ss→sz 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 sz→sz 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
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 uses computational models to infer three-dimensional protein shapes and complexes from sequences and other biological data.
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.
Kevin Lu, Jannik Brinkmann, Stefan Huber +4
1Northeastern University · 2TU Clausthal · 3Harvard University +2
Protein language models (pLMs) encode information about protein sequences which enable downstream tasks such as structure prediction, but their internal representations are not well understood. Sparse autoencoders (SAEs) provide a promising tool to disentangle latent pLM representations into interpretable features, but existing annotation pipelines largely rely on protein-level annotations derived from database labels and LLM annotations of top activating sequences. Such annotations can overlook the localized residue-level and geometric patterns encoded by sparse features. We introduce an automated and scalable method for interpreting SAE features in ESM-2 by using geometrically inspired features of the protein Cα backbone. Across ESM-2 8M layers, an FDR-controlled discovery analysis shows that local geometry is significantly associated with many SAE features, with varying levels of predictive strength, expanding coverage beyond database and sequence-based methods. In particular, geometry can distinguish SAE features sharing the same database annotation, revealing substructure within known biological labels. A significant portion of SAE features activate on unannotated metagenomic protein sequences enabling us to use our SAE annotations to better understand these sequences. In addition, ablation experiments at the level of contact prediction show that removing found geometric features shifts ESM-2's predicted contact maps in the direction of the descriptor. This provides a robust method of annotating proteins activated within SAE neurons at a residue level, providing a bridge between mechanistic interpretability and structural biology.
Foundation models for structural biology have achieved remarkable performance in predicting biomolecular structure and show promise for the design of proteins and small molecules. Yet understanding which internal features drive their outputs remains challenging. Standard sparse autoencoders (SAEs), effective on transformer-style sequence embeddings, do not transfer cleanly to pairformer-like architectures: naively operating on pairwise representations yields a quadratic blow-up of features and obscures concepts distributed jointly across sequence and pair representations. We introduce PairSAE, which summarizes pairwise tensors via an N-mode SVD into token-wise interaction roles, then uses a sparse autoencoder to learn a shared set of token-level features that decode into both sequence and pair representations. Evaluated on Boltz-2 activations for PLINDER protein-ligand complexes, PairSAE yields interpretable features that align with UniProt annotations and predict Boltz-2 affinity values. These results indicate that PairSAE links the latent space of foundation models for structural biology to interpretable structural concepts, clarifying what the model "knows" while avoiding pairformer-induced pitfalls that limit conventional SAEs.
Giosue Migliorini, Aristofanis Rontogiannis, Grigori Guitchounts +3
University of California, Irvine · Flagship Pioneering