PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Bohan Liu, Xiaohan Fan, Wenqian Ye, Aidong Zhang

arXiv:2610.36226v1q-bio.QMcs.AIcs.LG

TL;DR

Structure predictors remain difficult to explain at the level of their sequence-to-structure computation. The paper trains sparse dictionaries on ESMFold’s folding trunk, interprets their features, and traces their composition. The resulting account distinguishes inherited sequence semantics from trunk-computed geometry and connects concepts to fold quality and biochemical support.

  • Problem

    The computation that transforms amino-acid sequences into 3D structures remains opaque, limiting scientifically reliable explanations and verification against underlying biological processes.

    Source: §1 · p. 1

  • Method

    The authors train sparse dictionaries on both ESMFold trunk representations, interpret features with a reasoning language model, and trace their composition and relation to folding outputs.

    Source: §1 · p. 2

  • Results

    The pair representation is 91% geometric, trunk code adds up to +0.32 auroc on residue relations, and known biochemistry supports about 73% of circuit edges.

    Source: §5 · p. 7

  • Takeaways & Limitations

    The dictionaries provide an interpretable account in which the single representation retains sequence semantics and the pair representation computes geometry.

    Source: §5 · p. 7

Abstract

from arXiv · show

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.

1 Introduction

The paper targets the opaque computation that turns protein sequence into structure, interpreting ESMFold’s folding trunk to distinguish inherited sequence semantics from geometry computed within the trunk.

1 Introduction

  • The paper addresses the opacity of structure predictors, whose learned sequence-to-structure computation lacks scientifically reliable explanations or guarantees.The authors identify verification against scientific processes and linking hidden features to output structures as unresolved aims.Source: §1 · p. 1
  • The authors hypothesize that sequence semantics is inherited while geometry is computed in the trunk and materialized by the structure module.The trunk representations are supervised by the 3D frame-aligned structure loss and auxiliary read-out losses.Source: §1 · pp. 1–2, Figure 1 · p. 1
  • The work trains sparse-feature dictionaries on both trunk representations and interprets every feature with a reasoning language model, checking labels against quantitative context withheld from the prompt assembler.The analysis also traces feature composition across blocks and connects interpreted code to folding outputs.Source: §1 · p. 2
  • The trunk code adds up to +0.32 auroc on residue relations, while residue-level properties are already recoverable from the language-model states.The authors describe the division of labor as single-representation residue semantics and pair-representation geometry.Source: §1 · p. 2

2 Related Work

Prior work supplies tools for interpreting sequence-model features and protein representations, but this paper targets the folding trunk’s internal computation and pair representation.

2 Related Work

  • Sparse autoencoders decompose dense activations into overcomplete feature dictionaries, but reconstruction and selectivity alone do not establish causal use or semantic meaning.Interventional and correlational tools address these distinct questions; the paper uses a residue-identity control as a correlational test.Source: §2 · p. 2
  • Protein sequence-model embeddings encode structural and functional signals, while ESMFold’s trunk explicitly regresses from sequence embeddings to 3D geometry.Prior probing work finds secondary structure, contacts, and some functional properties linearly decodable from protein language-model embeddings.Source: §2 · p. 2
  • The paper applies sparse autoencoders to ESMFold’s single and folding-exclusive pair representations, then traces interpreted concepts across trunk blocks with integrated gradients.This addresses a gap left by sequence-backbone SAE work that related features to ESMFold outputs without interpreting the trunk’s internal computation.Source: §2 · pp. 2–3

3 Method

The method analyzes five ESMFold trunk blocks, trains sparse dictionaries on their single and pair states, interprets features, evaluates their contributions, and traces their composition across blocks.

3.1 The folding trunk and pilot analysis

  • Pair states encode pairwise residue distance at deep blocks, while single states contain little distance signal; functional-site signal erodes with depth as geometric signal such as burial builds.ESM-2–trunk CKA averages 0.14, compared with 0.65 among trunk blocks and 0.96 among middle blocks.Source: §3.1 · p. 3, Figure 2 · p. 4

3.2 Sparse dictionaries of the trunk

  • TopK sparse autoencoders are trained at five trunk blocks on single and pair states, with k=16 to balance reconstruction and sparsity.The models encode mean-centered residue or pair activations and train to minimize reconstruction error while keeping decoder rows at unit norm.Source: §3.2 · p. 3

3.3 Automated interpretation

  • A reasoning language model labels single-state features from residue contexts and pair-state features from populations of residue pairs and both endpoints’ contexts.The single-state prompt also assigns concept categories and asks whether a feature reflects amino-acid type or a biological concept.Source: §3.3 · p. 3

3.4 Feature evaluation

  • Feature evaluation localizes activation sites on folded structures and measures trunk-code contributions over sequence-model states using held-out probes.The protocol also probes code-based prediction of ESMFold confidence and ground-truth accuracy on unseen proteins.Source: §3.4 · pp. 3–4

3.5 Circuit construction by integrated gradients

  • Integrated gradients attribute each downstream feature to an upstream code through the live ESMFold pass, tracing within-track flow and coupling between single and pair representations.Cross-attribution distinguishes forward single-to-pair writes from backward pair-to-single writes.Source: §3.5 · p. 4

4 Experiments and Results

The experiments show that trunk codes add relational information beyond ESM-2, while interpreted single and pair features form distinct, partly coupled computations whose activations also track fold accuracy.

4.1 Implementation details

  • TopK dictionaries reconstruct single-representation activations across five trunk blocks with explained variance of 0.86–0.90 at k=16.Single-representation dictionaries use 95M activations from 500k folded sequences; pair dictionaries use 4.8M activations from 967 proteins, and the interpretation pipeline labels 5,595 single and 2,972 pair features.Source: §4.1 · pp. 4–5, Dictionary evaluation · p. 5

4.2 Growing relational information in the trunk

  • Adding trunk codes improves residue-contact AUROC by +0.318, ladder-continuation AUROC by +0.209, and strand-orientation AUROC by +0.086 beyond ESM-2-3B.For strand orientation, the trunk pair code reaches 0.965 versus 0.862 for ESM-2-3B; residue-level gains are generally within the shuffled control, and whole-protein retrieval matches ESM-2-3B.Source: Table 1 · p. 5, §4.2 · p. 5

4.3 Faithful automated interpretation of concepts

  • Pair-feature labels are 91% geometric, while single-feature labels are 57% geometric and become less so with depth, falling from 0.75 at b00 to 0.43 at b47.A cross-prompt control preserves the pair track’s geometric character, while independent checks find identity labels match the dominant residue in 12 of 13 clusters at 0.84 purity versus 0.40 for structural concepts; labeled geometry features also predict their corresponding relations.Source: §4.3 · pp. 5–6

4.4 Activation sites

  • For the selected detectors, disulfide activation concentrates on bonded cysteine pairs, while transmembrane, signal-peptide, and N-glycosylation detectors concentrate on their annotated spans; these detectors also annotate proteins lacking functional-site labels.Among cysteines flagged as disulfide-bonded across 201 such proteins, 594 of 1,121 take a Cβ–Cβ partner within 6.5 Å in the predicted fold.Source: §4.4 · p. 6

4.5 Direction of information flow

  • Integrated-gradient attributions align with known biochemistry for 79% of within-track edges, 75% of forward cross-track edges, and 63% of backward cross-track edges.Ablating an upstream feature moves its attributed downstream partner 1.42× as much as a magnitude- and direction-matched control; about one-third of pair-code attribution traces to the single code, versus about 6% in the reverse direction.Source: Figure 3 · p. 6, §4.5 · p. 7

4.6 Case study: biochemically supported and rejected edges

  • Among rejected cross-track attributions, 37% of backward edges are unsupported versus 25% of forward edges, and 248 of 290 rejected edges run backward.The selected rejected cases include local causes attributed to long-range geometry and residues linked to geometry against their chemical properties.Source: §4.6 · p. 7

4.7 Influence on the folding output

  • At read-out, interpretable-code predictions reach R2=0.84 for pLDDT and R2=0.60 for lDDT, up from R2=0.61 and R2=0.48 at the first block.Buried and packed concepts mark more accurate residues, while solvent-exposed, flexible, and terminal concepts mark less confident and less accurate ones.Source: Figure 5 · p. 7, §4.7 · p. 7

5 Conclusion

The trunk dictionaries reveal distinct roles for the single and pair representations and support a biochemistry-aligned circuit; supplementary figures show where interpreted detectors localize and how concepts relate across depth.

5 Conclusion

  • The pair representation is 91% geometric while the single preserves sequence semantics; trunk codes add up to +0.32 AUROC over language-model embeddings for residue-relation inference, and known biochemistry supports about 73% of circuit edges.Source: §5 · p. 7

A Supplementary interpretability figures

  • Disulfide and span-based detectors localize well, whereas ligand-binding, active-site, and modified-residue readouts are distributed and less accurate.Figure 6 paints per-residue detector maps and annotated sites across eight annotation types and four proteins each; green labels denote well-localized features and yellow labels distributed features.Source: A Supplementary interpretability figures · p. 10
  • Figure 7 profiles vocabulary geometry and label grounding across depth, while Figure 8 resolves category flow into individual concepts.Figure 8 encodes attributed mass as node area and summed weight as edge width, with both log-scaled.Source: A Supplementary interpretability figures · p. 10
  • The concept-flow rendering groups the vocabulary into representation-specific annotation categories and distinguishes forward single-to-pair from backward pair-to-single bundles.Every feature enters one category row, and the two cross-track bundles use a shared scale to encode their asymmetry.Source: A Supplementary interpretability figures · p. 10

B Controls and robustness

Control and robustness analyses test attributed feature composition, quantify output attribution, and specify uncertainty and sparsity choices. The reported interventions support downstream feature effects, while folding output influence remains distributed across features.

Controls and robustness

  • Confidence intervals use nonparametric percentile bootstraps, with evaluation units resampled by protein for detection and output-attribution analyses and by attributed target for causal fold-change ratios.The protocol reports the 2.5th and 97.5th percentiles as the 95% interval.Source: B Controls and robustness · p. 10
  • Trunk-feature permutation changes auroc by at most ±0.010 on every property, preserving feature marginals while removing alignment to residues.The relational increments are compared against this permutation band.Source: B Controls and robustness · p. 10

Feature composition

  • Read-out-matched ablations shift attributed downstream partners 1.42× the matched control in 62% of cases; pair-track tests yield 1.64× in 88%.The single-track intervention shifts distogram-L1 3.70× the matched control; the magnitude-only comparison is 56.6× and is described as the weaker test.Source: B Controls and robustness · p. 10

Output attribution

  • At b24, probes predict pLDDT and true per-residue accuracy at R^2=0.73 and 0.56, versus 0.17 and 0.08 for a representation-magnitude baseline; single-feature ablation leaves folding output consistent.ESMFold confidence and true accuracy agree at ρ=0.72, while output influence is distributed across multiple features.Source: B Controls and robustness · pp. 10, 12

Controls and robustness

  • Figure 6 compares per-residue detector maps with annotated sites across eight types, while Figure 7 profiles vocabulary geometry and label grounding across depth.The Figure 6 caption distinguishes well-localized from distributed features by green and yellow labels.Source: Figure 6 · p. 11, Figure 7 · p. 12
  • All main-text dictionaries use k=16, trading reconstruction for maximally interpretable codes despite reconstruction improving with larger active-feature budgets at every block.Table 3 reports held-out explained variance for the trunk single representation.Source: B Controls and robustness · p. 12, Table 3 · p. 12

C Amino-acid type as a measurement bias

The analysis controls for amino-acid identity because residue type can act both as a legitimate biological prior and as a measurement bias. Within-type scoring and selection reveal which detectors carry signal beyond that prior, with catalytic-site signal showing a boundary across structure sources.

Amino-acid type as a measurement bias

  • Amino-acid type is decodable at auroc=1.0, so detection is recomputed within protein and residue-type cells to measure feature signal beyond that prior.The control-task score is treated as attribution across the interpreted benchmark, not a pass/fail test.Source: C Amino-acid type as a measurement bias · p. 12
  • Holding residue type fixed gives within-AA auroc=0.87 for catalytic-site localization and 0.95 for burial; re-selection under this control yields glycosylation 0.782, active-site 0.974, and metal-binding 0.924.Raw-auroc selection gives glycosylation 0.950 raw but 0.563 within-AA; type-independent signal-peptide selection is unchanged.Source: C Amino-acid type as a measurement bias · p. 12
  • Table 4 reports residue-type-only auroc and dominant residues alongside annotated site types, while Table 5 describes detector choices, auroc, overlap, and permutation-baseline comparisons.Figure 8 categorizes single-representation concepts separately from pair-representation geometry classes.Source: Table 4 · p. 13, Table 5 · p. 13, Figure 8 · p. 13
  • Disulfide detection reaches auroc=0.97 at the deepest block, but a distance-to-centroid statistic already reaches 0.82 for catalytic sites on M-CSA.Applying the within-AA catalytic detector to single-sequence AlphaFold structures yields auroc between 0.51 and 0.62, consistent with the code riding evolutionary signal from ESM-2 L36 rather than trunk-computed geometry.Source: C Amino-acid type as a measurement bias · p. 13

D Applications of the interpreted detectors

The interpreted detectors identify residue-localized and distributed features, generate coordinate-level hypotheses for unannotated proteins, and localize sites in confidently folded disease proteins. Circuit attribution and representation probes distinguish trunk geometry from inherited sequence information.

D Applications of the interpreted detectors

  • Disulfide, metal-binding, and transmembrane features concentrate on annotated residues, significantly against a residue-permutation null on 68–91% of proteins.The general functional feature instead spreads across the active-site and ligand-binding pocket.Source: D Applications of the interpreted… · p. 13
  • Among 300 unannotated human proteins, 127 had at least one detector prediction that passed its respective spatial or sequence sanity check.Checks included predicted disulfide partners within 6.5 Å, signal residues in the N-terminal 35 residues, and glycosylation sites on the N-X-[ST] sequon.Source: D Applications of the interpreted… · pp. 13–14
  • On confident sites in 139 brain-disease proteins, active-site, metal-binding, and ligand-binding AUROCs were 0.899, 0.815, and 0.778, while evolutionary conservation remained competitive.Variant disruption reached auroc≈0.75, but per-feature ΔΔG prediction fell to auroc=0.55 under protein-grouped cross-validation.Source: D Applications of the interpreted… · p. 14

E Circuit attribution

  • Integrated gradients attribute downstream block codes to upstream codes by interpolating from a zero-code baseline and splicing decoded activations into the forward pass.The implementation uses midpoint Riemann sums and online gradient×input accumulation for constant peak memory in sequence length.Source: E Circuit attribution · p. 14

F Representational similarity between the sequence model and the trunk

  • Trunk blocks show stronger representational similarity to one another (CKA 0.65; middle blocks 0.96) than to ESM-2 (0.14 on average), supporting five-block sampling.ESM-2's final layer versus b00 reaches 0.23; the analysis samples 9,000 residues across 4,155 proteins.Source: F Representational similarity between… · p. 14

G Pairwise distance beyond ESM-2

  • Pair-representation dictionary codes predict Cα–Cα distance with R2=0.69 at b47 and 0.64 at b24, while single-residue probes, including the ESM-2-3B backbone, do not recover it.The comparison covers 60,000 residue pairs across 500 proteins and is scale-matched.Source: G Pairwise distance beyond ESM-2 · p. 14, Table 6 · p. 14

H Recycle similarity

The single representation is nearly unchanged between recycle budgets N=1 and N=3, and dictionaries select similar active features in both settings. Most of its final state is reached after the first recycle.

H Recycle similarity

  • At b00, b24, and b47, single representations from N=1 and N=3 agree at cosine 1.0000, 0.99997, and 0.99979, respectively.The dictionaries select the same active features on 87–98% of residues, with activation correlation 0.98–0.99; cosine similarity reaches at least 0.9998 of its final value by the first recycle.Source: H Recycle similarity · pp. 14–15

I Causal ablation on pair circuit

Ablating attributed pair features shifts their downstream partners and predicted geometry in the live model. Protein-level retrieval, however, shows no trunk-code advantage that persists across descriptor choices.

I Causal ablation on pair circuit

  • Across 1,129 edge–protein cases, pair-feature ablation moved the attributed downstream partner 1.64× a signed-cosine-matched control and shifted target-pair distogram-L1 1.81× the matched control.The partner movement exceeded the control on 88% of cases; against a magnitude-matched control alone, the ratio was 31.7×.Source: I Causal ablation on pair circuit · p. 15

J Protein-level retrieval

  • Protein-level retrieval showed no trunk-code advantage robust to pooling: outcomes ranged from +0.026 on enzyme class to −0.084 across tasks under alternative descriptors.The benchmark is dominated by sequence-homologous labeled proteins, while pooling discards the pairwise structure that distinguishes the trunk code.Source: J Protein-level retrieval · p. 15

K Pair-feature label evaluation

Pair-feature labels track distinct geometric relations, and swapping interpretation prompts preserves the geometric character of pair features while shifting single-feature labels away from geometry. Attribution summaries assign most mass to the pair representation.

K Pair-feature label evaluation

  • β-labeled features predict strand pairing at 0.938 versus helix contact at 0.839, while helix-labeled features predict helix contact at 0.989 versus strand pairing at 0.737.The double dissociation holds at both blocks and exceeds a random group of the same size.Source: K Pair-feature label evaluation · p. 15, Table 7 · p. 15

L Prompt inductive bias

  • Under swapped prompts, pair-feature labels remain 91–99% geometric, while single-feature geometry falls from 57% under its own question to 32% under the pair question.The authors interpret this as evidence that the representation distinction is not merely a prompt effect.Source: L Prompt inductive bias · pp. 15–16

M The full attribution

  • The full attribution assigns most attributed mass to the pair representation; the forward single-to-pair direction supplies about a third of it, and the backward direction is smallest.Figure 3 summarizes four attributions across the representation directions and block pairs.Source: M The full attribution · p. 16

N Biochemical sanity check

The biochemical rule set supports 79% of within-representation edges, while support differs between the forward and backward cross-track legs. Prompt formats and feature categories distinguish sequence-derived information from pairwise geometry, with some categories excluded from that split.

N Biochemical sanity check

  • The audit matches identity endpoints to residue keywords and structural endpoints to geometry classes, excluding catch-all pair geometry from support counts.Grounding includes same-residue identity matches, documented residue propensities, and established relations between structural categories.Source: N Biochemical sanity check · p. 16, Table 8 · p. 16
  • 79% of within-representation edges are supported: 609 of 770 in the single track and 120 of 152 in the pair track.The cross-track forward ss →sz leg reaches 75% of edges and 66% of attributed weight, versus 63% and 46% for the backward sz →ss leg.Source: N Biochemical sanity check · p. 16
  • Unsupported edges include locally acting features linked to long-range geometry, which the rule set rejects because it lacks a recorded relation.The authors describe those rejections as reflecting limited rule-set coverage, while other unsupported edges pair residues with geometry their chemistry resists; propensities are statistical and further verification is needed.Source: N Biochemical sanity check · p. 16

O Auto-interpretation prompt details

  • The two representations use different feature-context formats and questions; a control swaps the questions while keeping each context block fixed.The prompts are reproduced in Figures 13 and 14.Source: O Auto-interpretation prompt details · p. 16

P Splitting the vocabulary into sequence and geometry

  • The vocabulary split counts residue identity and sequence offset as sequence-determined, and residue relations or 3D environment as geometric.Secondary structure, functional site, and disulfide are excluded from both sides; the reported shares use only features whose category settles the classification.Source: P Splitting the vocabulary into… · pp. 16, 20

Related papers

Loading 2610.36226v1…