Research map / Probabilistic and scientific modeling / Molecules and proteins
From Mechanistic Interpretability to Mechanistic Biology: Training, Evaluating, and Interpreting Sparse Autoencoders on Protein Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Trains TopK sparse autoencoders on residual-stream activations of ESM-2 (650M), characterizes the latents (family specificity, activation patterns), and uses linear probes on SAE latents to link interpretable features to known protein determinants.
Problem
Protein language models are powerful but their learned features are poorly understood; it is unclear what biological concepts pLMs encode and whether those representations can reveal novel protein biology.
Approach
Apply TopK sparse autoencoders (mechanistic interpretability) to ESM-2 layers, classify latents by activation pattern and protein-family specificity, and interpret linear probes built on SAE latents to map predictive features to known or candidate biological mechanisms.
Key results
- Approximately 80% of SAE latents were rated interpretable by blinded human raters, versus far fewer ESM neuron latents from layer 24.
- Linear probes on SAE latents perform comparably to ESM embedding probes across layers on four tasks; for secondary structure SAE probes consistently outperform ESM in most layers.
- Lower k (higher sparsity) and larger expansion factor both increase the number of family-specific latents; family-specific latents peak in early-to-mid layers.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train a TopK SAE on ESM-2 residual-stream activations instead of using raw ESM latents | ESM-2 layer 24 · Manual interpretability study | interpretable ratings: Approximately 80% | Helped |
| Train linear probes on SAE embeddings instead of ESM embeddings | ESM-2 · TAPE secondary structure | accuracy | Helped |
| Train linear probes on SAE embeddings instead of ESM embeddings | ESM-2 · Subcellular localization | accuracy | No clear effect |
| Train linear probes on SAE embeddings instead of ESM embeddings | ESM-2 · Melteome Atlas thermostability | Spearman's rank correlation | No clear effect |
| Train linear probes on SAE embeddings instead of ESM embeddings | ESM-2 · CHO cell expression | accuracy | No clear effect |
| Increase TopK sparsity by lowering k (ablation) | TopK SAE · Swiss-Prot protein-family classification | number of family-specific latents | Helped |
| Increase the SAE expansion factor (ablation) | TopK SAE · Swiss-Prot protein-family classification | number of family-specific latents | Helped |
| Increase the SAE expansion factor (ablation) | TopK SAE · Latent activation-pattern classification | activation-pattern distribution | No clear effect |
| Use middle ESM layers rather than later layers (ablation) | TopK SAE on ESM-2 · Swiss-Prot protein-family classification | number of family-specific latents | Helped |
| Steer a family-specific latent rather than a random active latent | ESM-2 · Protein sequence steering | number of residues changed | Helped |
Surprising findings
- Thermostability probes rely on simple amino acid composition latents (e.g., glutamine), and middle layers performed worse than early or late layers for this task.
- Some highly predictive latents, such as a membrane localization latent, activate predominantly on poly-alanine sequences with unclear function.
Limitations
- Only the 650M-parameter ESM-2 variant was studied; other sizes, checkpoints, and architectures were not analyzed.
- Feature sets depend on the training data distribution, and mean-pooling for protein-level probes has known shortcomings.
- Latent interpretations are constrained by current biological knowledge and may be incomplete.
Experimental setup
Datasets: UniRef50 TAPE secondary structure dataset Almagro Armenteros subcellular localization dataset Melteome Atlas thermostability dataset (via FLIP) Masson et al. 2022 CHO expression dataset Swiss-Prot InterPro protein families
Benchmarks: TAPE secondary structure subcellular localization thermostability (FLIP split) mammalian CHO cell expression
Metrics: accuracy Spearman rank correlation F1 score human interpretability yes-rate
Baselines: ESM-2 embedding linear probes Simon & Zou 2024 SAE on ESM-2
Builds on: TopK sparse autoencoders (Gao et al., 2024) Towards monosemanticity / sparse dictionary learning (Bricken et al., 2024) Li et al. 2024 pLM downstream-task analysis TAPE benchmark
Closest papers by meaning
- Feature Reuse and Scaling: Understanding Transfer Learning with Protein Language Models (ICML 2024 · similarity 0.90)
- Training Compute-Optimal Protein Language Models (NeurIPS 2024 · similarity 0.89)
- SaProt: Protein Language Modeling with Structure-aware Vocabulary (ICLR 2024 · similarity 0.89)
- Metalic: Meta-Learning In-Context with Protein Language Models (ICLR 2025 · similarity 0.88)
- LEMON-ZEST: Evolution-Informed Tokenization for Efficient Protein Language Modeling (NeurIPS 2026 · similarity 0.88)
- Understanding protein function with a multimodal retrieval-augmented foundation model (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.