Opening the Black Box of Protein Folding: Sparse Autoencoders Meet ESMFold

Protein language models have become the workhorses of computational structural biology. Models like ESM2 and the structure prediction systems built on top of them, such as ESMFold, can map an amino acid sequence to a three dimensional structure with accuracy that rivals experimental methods on many targets. Yet the internal logic of these models remains almost entirely opaque. We know they work, but we do not know how the network decides that a particular stretch of residues forms an alpha helix, or that two distant positions in the sequence will end up in contact. The paper "Towards Interpretable Protein Structure Prediction with Sparse Autoencoders" by Parsan, Yang, and Yang, presented at the GEM workshop at ICLR 2025, takes a serious step toward changing that. Below I unpack what they did, why the methodological choices matter, and where I think this line of work is heading.

Why Interpretability of PLMs Is Hard

The dominant narrative in mechanistic interpretability of protein models has been attention analysis. Rao et al. and Vig et al. showed that attention maps in transformer PLMs encode residue-residue contacts, and Zhang et al. argued that PLMs predict structure largely by memorizing coevolutionary patterns. These are valuable findings, but they share a fundamental limitation: they are correlational. Observing that an attention head lights up when two residues are in contact does not tell you whether that head is causally responsible for the prediction, or whether it is merely a downstream reflection of computation happening elsewhere.

Sparse autoencoders offer a different angle. The core idea, borrowed from interpretability work on large language models (Templeton et al., Gao et al.), is that the activations of a neural network are dense and polysemantic: any single neuron participates in representing many unrelated concepts. An SAE learns to decompose these dense activations into a much higher dimensional but sparse set of latent features, each of which ideally corresponds to a single human-interpretable concept. Formally, given token embeddings x in R^d, the SAE encodes them into z in R^n with d << n, enforces sparsity on z, and decodes back to a reconstruction x_hat by minimizing L2 reconstruction loss.

Prior work (Simon and Zou, Adams et al.) applied SAEs to smaller ESM2 variants and found thousands of biologically interpretable features. The catch is that those smaller models cannot actually do structure prediction. ESMFold is built on ESM2-3B, and nobody had scaled SAEs to that model. This is the first gap the paper closes.

Two Technical Advances

The contribution has two pillars worth examining separately.

Scaling SAEs to ESM2-3B. This is more than an engineering footnote. By training SAEs on layers 18 and 36 of the 3 billion parameter model, the authors can finally intervene at the representation that actually feeds ESMFold's structure trunk. Their training set is substantial: 10 million sequences from UniRef50, roughly 2.5 billion tokens. They normalize embeddings during training to allow hyperparameter transfer between layers, a sensible practical choice that reduces the tuning burden across the network's depth.

Matryoshka SAEs for hierarchical features. This is the part I find most conceptually elegant. A standard SAE has no notion of structure among its latents; all features sit in one flat dictionary. The Matryoshka variant divides the latent dictionary into nested groups, where each nested prefix must reconstruct the input independently. The smallest group has to do a passable job on its own, then progressively larger groups refine the reconstruction. The effect is a forced coarse-to-fine hierarchy: early groups learn abstract, broadly useful features while later groups encode granular detail.

The motivation is biological. Proteins are organized across scales, from local backbone motifs and secondary structure elements up through domains and full molecular assemblies. A flat dictionary has no inductive bias toward this organization; the Matryoshka structure imposes one. Whether the learned hierarchy genuinely aligns with these biological scales is an empirical question, but the prior is well motivated.

What the Evaluations Show

The authors are careful to demonstrate that their SAEs do not destroy the very capabilities they want to study. Two evaluation axes matter here.

On language modeling, they measure the change in cross-entropy loss between the original PLM logits and the logits after passing through the SAE reconstruction. At low sparsity (L0 below 10), all three architectures (TopK, Matryoshka, and L1) degrade similarly, with delta CE around 2.5 to 3.0. As sparsity budget increases, TopK and Matryoshka pull ahead of L1, with delta CE approaching zero above L0 of 100. The takeaway is that Matryoshka matches or beats the standard architectures, so the hierarchical structure comes at no reconstruction cost.

On structure prediction, the results are more striking. Because ESMFold consumes representations from all layers but the SAE reconstructs only one, the authors ablate all other layers to isolate the effect. The backbone RMSD comparisons on CASP14 are telling:

• Experimental vs. ESMFold baseline: 3.1 angstroms
• Experimental vs. ESMFold with full ablation: 15.1 angstroms
• Experimental vs. ESMFold using only layer 36: 2.9 angstroms
• Experimental vs. SAE reconstruction of layer 36: 3.2 angstroms

The full ablation collapses prediction quality to 15.1 angstroms, confirming the layers carry real information. But keeping only layer 36, or replacing it with the SAE reconstruction, preserves accuracy at roughly baseline levels. This is a clean demonstration that layer 36 alone carries most of the structurally relevant signal, and that the SAE captures it faithfully.

The Steering Result and Why It Matters

The most important conceptual move is the steering experiment. By intervening on the layer 36 representation while holding the input sequence fixed, the authors increased the solvent accessibility of the predicted structure. This is causal interpretability in the strict sense: they identified features, manipulated them, and observed a directed, interpretable change in the output.

This is qualitatively different from attention map analysis. An attention map tells you what correlates with contacts. A steering intervention tells you that a specific feature causes a structural property to shift. The distinction matters enormously for downstream applications. If you can identify a latent that controls a structural property and turn it up or down at will, you have the beginnings of a controllable design tool, not just an analysis instrument.

My Take

What excites me intellectually about this work is that it reframes ESMFold as a system you can perform surgery on, rather than a sealed oracle. A few observations and caveats:

First, the finding that layer 36 alone nearly matches the full model is more significant than the paper emphasizes. It suggests substantial redundancy in how ESMFold consumes ESM2's representations, and it identifies a single privileged layer as the natural intervention point. For anyone designing follow-up interpretability or control experiments, this is a useful concentration of signal.

Second, the steering demonstration is genuinely a proof of concept and should be read as such. Increasing solvent accessibility is a relatively coarse, scalar property. The harder and more valuable question is whether one can steer toward specific tertiary arrangements, particular fold families, or designed binding interfaces while keeping the sequence fixed. Solvent accessibility is the right starting point precisely because it is measurable and continuous, but the design relevance grows sharply with the specificity of the target property.

Third, I would want to see whether the Matryoshka hierarchy maps onto the biological scale hierarchy in a verifiable way. The architecture encourages coarse-to-fine organization, but the paper's evidence that this aligns with secondary-to-tertiary structure is still early. A convincing demonstration would show that pruning to the smallest nested group recovers secondary structure features while larger groups add domain-level features.

Limitations and Open Questions

The honest limitations are worth stating. The structure evaluation relies on layer ablation, which is a somewhat artificial intervention; the real ESMFold uses all layers jointly, and isolating one layer may not reflect the full computational picture. CASP14, while challenging, is a limited test set. And steering a single scalar property does not yet establish that fine-grained, multi-feature control is feasible without introducing physically implausible structures.

The open questions I find most compelling: Can these features be composed to design novel folds? Do the discovered latents transfer across model scales and architectures, or are they idiosyncratic to ESM2-3B? And can causal feature attribution eventually replace coevolution as our mechanistic story for how PLMs fold proteins?

The authors open-sourced their code, dataset, pretrained models, and a visualizer, which lowers the barrier for the community to pursue these questions. That, more than any single result, is what makes this paper a useful foundation. We now have a concrete handle on the internals of a working structure predictor, and the interesting work of turning that handle into design capability is just beginning.