SAEFold

Protein language models now sit underneath most sequence-to-structure pipelines, but their internal representations are dense and entangled: a folding prediction emerges from a high-dimensional embedding with no obvious mapping back to biology. Sparse autoencoders (SAEs) attack this by learning an overcomplete dictionary in which each embedding is reconstructed from a handful of active latents, so individual directions can be read as concepts. Until SAEFold, that technique had only been applied to protein language models small enough to be interpretable but too small to fold anything.

SAEFold closes that gap by training SAEs on ESM2-3B, the encoder that ESMFold builds its structure module on top of. Two changes make this work. The first is simply scale: SAEs trained on the 3B model recover far more biological concepts than those trained on smaller checkpoints, and their reconstructions preserve enough signal for the downstream folding head to still produce accurate structures. The second is the adaptation of Matryoshka SAEs to protein representations, where nested groups of latents are each required to reconstruct the input independently, producing a coarse-to-fine hierarchy of features rather than a flat dictionary.

The work was produced by Reticular, a San Francisco startup working on interpretable models for protein engineering, with a collaborator at the University of Pennsylvania, and was presented at the ICLR 2025 GEM workshop. It sits alongside other sparse-dictionary efforts in biology such as PLM-SAE for variant effect prediction and CytoSAE for microscopy embeddings, but is the one aimed squarely at the folding pathway.

#Key Features

• Interpretability at folding scale: SAEs trained on ESM2-3B representations attach human-readable features to the exact embeddings ESMFold consumes, so a prediction can be attributed to concepts rather than treated as a black box.
• Matryoshka hierarchy: Nested latent groups each reconstruct the input on their own, so a small group captures broad structure while larger groups add detail, matching or beating standard TopK SAEs on reconstruction quality.
• Structure-preserving reconstruction: Replacing the language model embedding with its sparse reconstruction leaves ESMFold's accuracy essentially intact, which is what makes feature-level analysis meaningful rather than a lossy approximation.
• Targeted steering: Amplifying or suppressing a single feature shifts the predicted structure in a controlled direction while the input sequence is held constant.
• Released artifacts: Training and evaluation code, pretrained SAE checkpoints, and an interactive feature visualizer are published under Apache 2.0.

#Technical Details

SAEs were trained on activations from layers 18 and 36 of ESM2-3B using 10 million UniRef50 sequences (roughly 2.5 billion tokens, truncated at 1,022 residues), with a dictionary of 20,480 latents — an 8× expansion over the model's hidden width — and sparsity levels from k=20 to k=100 active latents per token. On Swiss-Prot concept discovery, the ESM2-3B Matryoshka SAE matched 233 annotated concepts at F1 above 0.5 (48.9% of those tested) across 2,677 feature-concept pairs, against 72 to 95 concepts for SAEs trained on ESM2-8M. Contact-map precision at L/2 was likewise consistently higher for the 3B-trained dictionaries. On CASP14 targets, baseline ESMFold predictions sat at 3.1 ± 2.5 Å RMSD to experiment and SAE-reconstructed predictions at 3.2 ± 2.6 Å, with as few as 8 to 32 active latents per token recovering most of the structure prediction signal. In the steering case study, applying a coefficient of −0.275 to a hydrophobicity feature raised total solvent-accessible surface area by 31.5%, from 8,369.5 to 11,009.3 Ų, while the backbone moved only 2.76 Å RMSD.

#Applications

The immediate audience is researchers who need to know why a structure predictor produced a given fold — debugging confident-but-wrong predictions, auditing model behavior on engineered or non-natural sequences, and identifying which learned concepts a protein family activates. The steering result points at a second use: biasing a predicted structure toward a desired physical property, such as greater surface exposure, without editing the sequence, which is directly relevant to solubility and interface engineering workflows. The public feature visualizer lets protein scientists inspect individual latents against real sequences without training anything themselves.

#Impact

SAEFold established that sparse dictionary learning transfers to the largest protein language models in routine use, and that the resulting features are good enough to both explain and intervene on structure prediction. Its Apache 2.0 release of code, pretrained dictionaries, and a hosted visualizer lowers the barrier for follow-on interpretability work considerably. The limits are worth stating plainly: the steering result is a single in-silico case study with no experimental validation, concept coverage tops out below half of the Swiss-Prot annotations tested, and the work appeared as a workshop paper and preprint rather than in a peer-reviewed venue.