PaperlayerView exact revision on arXiv

Source-linked AI summary

InterPLM: Discovering Interpretable Features in Protein Language Models via Sparse Autoencoders

Elana Simon, James Zou

arXiv:2412.12101v1q-bio.BMcs.AIcs.LGq-bio.QM

TL;DR

Protein language models achieve strong structure and function prediction, yet their learned representations remain difficult to interpret. This paper applies sparse autoencoders to ESM-2 embeddings and finds many concept-aligned features compared with individual neurons. The features also support annotation discovery and targeted steering, though steering more complex biological concepts remains challenging.

  • Problem

    Protein language models predict structure and function successfully, but their internal mechanisms and learned representations remain poorly understood.

    Source: §1 · p. 1

  • Method

    The paper trains sparse autoencoders on ESM-2 embeddings and combines feature analysis with automated language-model descriptions and biological annotations.

    Source: §1 · p. 3, §3.5 · p. 8

  • Results

    Up to 2,548 SAE features per layer had strong concept alignment, versus up to 46 clear concept associations per layer for neurons; SAE features covered 143 concepts versus 15 for neurons.

    Source: §3.3 · p. 6

  • Takeaways & Limitations

    Feature activation patterns can suggest missing protein-database annotations and support targeted control of model predictions through interpretable steering.

    Source: §4 · p. 10

    Steering was demonstrated for periodic glycine features in fairly simple contexts; extending it to complex biological concepts requires robust methods and clear success metrics.

    Source: §4.3 · p. 11

Abstract

from arXiv · show

Protein language models (PLMs) have demonstrated remarkable success in protein modeling and design, yet their internal mechanisms for predicting structure and function remain poorly understood. Here we present a systematic approach to extract and analyze interpretable features from PLMs using sparse autoencoders (SAEs). By training SAEs on embeddings from the PLM ESM-2, we identify up to 2,548 human-interpretable latent features per layer that strongly correlate with up to 143 known biological concepts such as binding sites, structural motifs, and functional domains. In contrast, examining individual neurons in ESM-2 reveals up to 46 neurons per layer with clear conceptual alignment across 15 known concepts, suggesting that PLMs represent most concepts in superposition. Beyond capturing known annotations, we show that ESM-2 learns coherent concepts that do not map onto existing annotations and propose a pipeline using language models to automatically interpret novel latent features learned by the SAEs. As practical applications, we demonstrate how these latent features can fill in missing annotations in protein databases and enable targeted steering of protein sequence generation. Our results demonstrate that PLMs encode rich, interpretable representations of protein biology and we propose a systematic framework to extract and analyze these latent features. In the process, we recover both known biology and potentially new protein motifs. As community resources, we introduce InterPLM (interPLM.ai), an interactive visualization platform for exploring and analyzing learned PLM features, and release code for training and analysis at github.com/ElanaPearl/interPLM.

1 Introduction

Protein language models predict protein structure and function successfully, but what they learn internally remains poorly understood. This work uses sparse autoencoders to interpret their learned features and connect them to biological concepts and applications.

1 Introduction

  • Understanding PLM mechanisms can inform model design, reveal biases and generalizability, and suggest biological insights beyond performance metrics.Interpretation may also help assess whether models learn physical and chemical principles or memorize structural motifs.Source: §1 · p. 1
  • Figure 1 diagrams embedding extraction, feature conversion, and reinsertion of reconstructed embeddings, alongside examples of structural and concept-associated activation patterns.The examples visualize activation intensity along protein sequences and structures.Source: Figure 1 · p. 2
  • The framework uses sparse autoencoders to analyze ESM-2 embeddings, characterize biological features, identify missing or new annotations, and steer model outputs.It also introduces an interactive feature-exploration platform and releases training and analysis code.Source: §1 · p. 3

2 Related Works & Background

Prior PLM interpretation work motivates questions about how models use conserved motifs, while sparse autoencoders offer a way to study feature-level representations. This paper frames biological feature interpretation as an evaluation challenge and presents a framework for analyzing PLM features.

2.1 PLM Interpretability

  • Even if PLMs memorize conserved evolutionary motifs, how they detect them in individual sequences and use them for prediction remains open.The questions also include what share of learned features focus on conserved motifs and what other computational strategies support predictions.Source: §2.1 · p. 3

2.2 Sparse Autoencoders (SAEs)

  • Neurons may encode unrelated concepts in superposition [10]; SAEs expand activations into sparse features that reconstruct them [12] [13].Each feature combines a dictionary vector with an activation value.Source: §2.2 · p. 3
  • Because biological feature interpretation requires domain expertise, the paper presents quantitative metrics and methods to visualize, analyze, describe, steer, and learn from PLM-derived SAE features.Interpretability is harder to assess than reconstruction quality, which has straightforward technical metrics.Source: §2.2 · p. 4

3 Results

SAEs reveal many more concept-aligned features than individual ESM-2 neurons, while feature clustering, automated descriptions, and targeted interventions expose interpretable biological patterns. The results include examples of missing annotation detection and steering glycine predictions beyond directly modified residues.

3.1 Sparse Autoencoders Find Interpretable Concepts in Protein Language Models

  • SAE features expanded ESM-2-8M's six layers from 320 neurons to 10,420 latent features, showing structural, protein-wide, and functional amino-acid recognition patterns.The patterns include coordinated activation among nearby residues, broad activation across proteins, and consistent activation in annotated domains, motifs, and binding sites.Source: §3.1 · p. 4

3.2 InterPLM: An Interactive Dashboard for Feature Exploration

  • InterPLM provides complementary views of feature activation, protein coverage, UMAP similarity, and Swiss-Prot alignment, with patterns often consistent across diverse proteins.Its views span features responding to single amino acids across many proteins and features responding to complete protein-family sequences.Source: Figure 2 · p. 4, §3.2 · p. 5

3.3 Sparse Autoencoder Features Capture More Biological Annotations than Neurons

  • Up to 2,548 SAE features per layer align strongly with 143 Swiss-Prot concepts, versus at most 46 neurons across 15 concepts; randomized-weight controls lose biological concept associations.The difference persists across layers; randomized models still yield features associated with amino-acid types.Source: Figure 3 · p. 5, §3.3 · pp. 6–7

3.4 Features Form Clusters Based on Shared Functional and Structural Roles

  • Feature clusters capture related roles with specializations, from kinase-binding-site preferences to specific versus general beta-barrel detectors.The TBDR detector f/1503 has F1=.998, while broader beta-barrel features have F1 scores .793 and .611; a separate cluster includes unlabeled features identified as glycosyltransferases.Source: Figure 4 · p. 6, §3.4 · p. 7

3.5 Large Language Models Can Generate Meaningful Feature Descriptions

  • Claude-generated descriptions predicted feature activations on separate proteins with median Pearson r=0.72, including motifs not represented by Swiss-Prot annotations.Description accuracy had only weak dependence on Swiss-Prot concept prediction performance (Pearson r=0.11).Source: §3.5 · p. 8

3.6 Feature Activations Identify Missing and Novel Protein Annotations

  • Feature activations identify likely missing Swiss-Prot annotations for Nudix motifs and peptidase S1 domains, while suggesting UDP and Mg2+ binding sites in glycosyltransferases.InterPro confirms a Nudix motif and S1 domain in example proteins; proposed binding-site annotations for unlabeled proteins require further validation.Source: §3.6 · pp. 8–9, Figure 6 · p. 8

3.7 Protein Sequence Generation Can be Steered by Activating Interpretable Features

  • Activating a periodic glycine feature at one residue increased glycine predictions at that position and the masked position three residues later, with effects propagating across repeats.Highly glycine-specific features influenced only the modified position, whereas periodic features produced diminishing effects across subsequent repeats; the scope across contexts and more complex patterns remains unresolved.Source: §3.7 · pp. 9–10, Figure 7 · p. 9

4 Discussion

The discussion argues that SAE features expose biological patterns represented in superposition and support annotation, model-analysis, and steering applications. It also identifies open questions about feature interpretation, evaluation, and extending steering to complex biological outcomes.

4 Discussion

  • SAEs extract thousands of interpretable features, and their activation patterns can reveal missing database annotations and support targeted control of sequence generation.The framework combines large-language-model interpretation with biological database annotations for scalable feature characterization.Source: §4 · p. 10

4.1 Implications for PLM Interpretability

  • ESM-2 yielded 2,548 SAE features per layer versus 46 interpretable neurons per layer against Swiss-Prot concepts, supporting feature superposition in PLM representations.SAEs trained on randomized PLMs captured individual amino acids but not the tested complex biological concepts, suggesting meaningful feature extraction requires learned weights.Source: §4.1 · p. 10
  • Domain-level F1 scores provide a quantitative way to assess feature-concept associations when annotations are coarser than learned features.The authors caution that these metrics approximate biological interpretability, while distinguishing approaches and potentially enabling systematic comparisons.Source: §4.1 · p. 10
  • The share of features devoted to motif memorization versus other computational strategies remains unclear, and the mechanisms linking features to masked-token predictions need further study.The framework may help quantify the motif balance, extending work on PLMs’ coevolutionary statistics.Source: §4.1 · p. 10

4.2 Applications of Interpretable PLM Features

  • Tracking features during training or fine-tuning could illuminate knowledge acquisition, while analyzing failure-mode activations could highlight systematic biases and guide model improvements.The framework could also support comparisons of learning patterns across PLMs and reveal how architectural choices influence biological concept capture.Source: §4.2 · p. 10
  • Consistent feature activations across protein families that lack matching annotations may point to unrecognized biological motifs or relationships.The authors also demonstrate targeted influence on sequence generation through feature-based steering, while describing periodic glycine steering in simple contexts as a promising direction rather than a transformative engineering result.Source: §4.2 · p. 10

4.3 Limitations and Future Directions

  • Extending steering to complex biological concepts remains challenging because robust methods and clear success metrics are needed.Validating intended biological changes without disrupting other properties is more complex than evaluating semantic shifts in language models.Source: §4.3 · p. 11
  • The analysis covers unmasked amino-acid representations, leaving masked-token and CLS representations unexamined, and mapping features into interpretable circuits remains open.The authors also identify systematic characterization of uncharacterized features as a possible route to discovering biological motifs.Source: §4.3 · p. 11

5 Methods

The methods train and evaluate ESM-2 sparse autoencoders against Swiss-Prot concepts, compare them with randomized-model and neuron baselines, and analyze feature descriptions, clustering, structure, and steering.

5.1.1 Dataset Preparation

  • The training data comprised 5M UniRef50 sequences, with ESM-2-8M representations taken after layers 1–6 and <cls> and <eos> tokens excluded.Sequences were sharded by 1,000 proteins and tokens shuffled within groups for random sampling.Source: §5.1.1 · p. 11

5.1.2 Architecture and Training Parameters

  • Each layer used 20 SAEs with 10,240 dictionary features for 320-dimensional embeddings, trained for 500,000 steps with sampled learning rates and L1 penalties.The expansion factor was 32x; learning rates and L1 penalties were ramped from zero during the initial training phase.Source: §5.1.2 · p. 11

5.1.3 Feature Normalization

  • Feature activations were normalized using each feature’s maximum on 50,000 Swiss-Prot proteins, with reciprocal encoder and decoder rescaling preserving reconstruction values.Source: §5.1.3 · p. 11

5.2.1 Dataset Construction

  • The concept-evaluation set contained 50,000 Swiss-Prot proteins split into 25,000-protein validation and test sets, retaining sufficiently prevalent concepts.Protein-level annotations were converted to binary amino-acid-level labels while retaining domain-level relationships for multi-residue annotations.Source: §5.2.1 · p. 11

5.2.2 Feature-Concept Association Analysis

  • Feature-on/off labels used activation thresholds of 0, 0.15, 0.5, 0.6, and 0.8, and feature-concept pairs were evaluated with modified precision, recall, and F1.The threshold producing the highest F1 score was selected for each feature-concept pair.Source: §5.2.2 · p. 12

5.2.3 Model Selection and Evaluation

  • The best model per layer was selected by average top feature-concept F1 on a validation subset, then test metrics evaluated selected features and validation pairs above F1 0.5.Initial selection used 20% of the validation set and 135 concepts; the test procedure reported selected-feature F1 per concept and how many validation pairs above 0.5 also exceeded 0.5 on test.Source: §5.2.3 · p. 12

5.2.4 Baselines; 5.5 Steering Experiments

  • Baselines used shuffled ESM-2 weights and biases or scaled neuron activations represented by an identity SAE; steering clamped selected features in SAE reconstructions before model predictions continued.The steering procedure recombined modified reconstructions with error terms and compared softmax probabilities across conditions.Source: §5.2.4 · p. 12, §5.5 · p. 13

5.3.1 Example Selection; 5.3.2 Description Generation and Validation; 5.4.1 UMAP Embedding and Clustering; 5.4.2 Sequential and Structural Feature Analysis

  • Feature analyses sampled high-activation proteins for language-model descriptions and validation, embedded decoder weights with UMAP, and assessed sequence and structural clustering against permutations.The structural analysis used proteins with AlphaFold structures, tested activation neighborhoods against random permutations, and excluded features with too few examples; descriptions were validated by comparing predicted and measured activation with Pearson correlation.Source: §5.3.1 · p. 12, §5.3.2 · p. 12, §5.4.1 · p. 13

B Extended SAE Details and Analysis

The extended SAE background describes sparse latent activations and dictionary vectors as components of feature analysis, and points to an overview of decomposition and training.

B.1 Additional background

  • An SAE expands a single amino acid’s latent vector into a larger, sparse vector whose nonzero positions indicate patterns in the original embedding.The paper analyzes feature activation values separately from learned dictionary weights.Source: §B.1 · p. 14
  • Figure 8 depicts an embedding as a weighted sum of dictionary elements and presents the SAE architecture.Source: Figure 8 · p. 14
  • The paper uses one SAE architecture while noting that alternatives vary initialization, nonlinearities, loss functions, and other training details to balance sparsity and reconstruction accuracy.Source: §B.1 · p. 14

B.1.1 SAE metrics

The methods background frames SAE training as a balance between sparse features and faithful reconstruction, with sparsity and recovered model loss among the evaluation measures.

B.1.1 SAE metrics

  • SAE training optimizes a weighted combination of an L1 sparsity term and mean squared reconstruction error because interpretability and accuracy are difficult to optimize directly.Source: §B.1.1 · p. 14
  • L0 measures the average number of nonzero latent values per token, while Percent Loss Recovered compares reconstructed-embedding cross-entropy with original and zeroed-layer cross-entropy.Source: §B.1.1 · p. 14
  • Final model selection used Swiss-Prot concept associations, and Table 1 is identified as reporting layer-wise learning parameters and SAE performance metrics.The authors state that L0 and Percent Loss Recovered informed initial hyperparameter-range experiments.Source: §B.1.1 · p. 14, Table 1 · p. 15

C Swiss-Prot Concept Associations

The Swiss-Prot metadata tables organize fields used for quantitative feature-concept associations and LLM feature descriptions and validation.

C.1 Swiss-Prot Metadata Categories

  • The C.1 tables group Swiss-Prot metadata fields by theme and mark whether each field supports quantitative concept analysis, LLM descriptions, or both.These fields underpin quantitative feature-concept associations and LLM feature descriptions and validation.Source: §C.1 · p. 15

C.2 Swiss-Prot Metadata Categories

  • The C.2 tables likewise group Swiss-Prot metadata by theme and indicate whether each field is used for quantitative analysis, LLM descriptions, or both.The metadata supports feature-concept associations and LLM description and validation.Source: §C.2 · p. 15

C.2.6 Functional Annotations

This appendix lists functional-annotation concepts associated with SAE features and describes prompts for generating feature summaries and predicting activation levels.

C.2.6 Functional Annotations

  • The functional-annotation examples include acyl-ester, DNA-tyrosine, and phosphohistidine intermediates, plus pantetheine-phosphoryl serine.Tables 8 and 9 list Swiss-Prot concepts associated with SAE features in any layer; an asterisk marks concepts also associated with a neuron.Source: §C.2.6 · p. 17, Table 8 · p. 17, Table 9 · p. 18
  • Figure 9 presents additional feature-concept analysis across layers.Source: Figure 9 · p. 18

D.1 Prompts

  • The description prompt asks an LLM to summarize metadata patterns associated with activation and identify sequence positions, nearby residues, and relevant annotations.It requests a concise characterization of the biological feature or pattern detected by the feature.Source: §D.1 · p. 19
  • The prediction prompt asks the LLM to estimate each query protein’s maximum feature activation from its metadata and the feature description.Predictions range from 0.0 to 1.0 and should reflect how well the protein matches the described pattern.Source: §D.1 · p. 19

D.2 Example Descriptions and Summaries

The example description characterizes feature 9047 as activating strongly on glycosyltransferases, with activation concentrated at sugar-binding regions.

D.2 Example Descriptions and Summaries

  • Highest activations of 0.9–1.0 occur in glycosyltransferases, especially glycogen synthases, around residues 280–450 at glycine, alanine, or proline sites.These sites frequently overlap nucleotide-sugar binding regions; related transferases and synthetases show medium activation, while proteins without sugar or nucleotide binding functions show none.Source: §D.2 · p. 20
  • Table 10 compares Pearson correlations for example feature-description summaries used alongside more verbose descriptions to predict maximum activation levels.Source: Table 10 · p. 20

E Additional cluster example

The glycosyltransferase cluster contains features with a shared theme but differing protein selectivity and activation locations.

E Additional cluster example

  • Figure 10 reports a small correlation between Swiss-Prot concept-label quality and the accuracy of LLM-predicted feature activation patterns.Source: Figure 10 · p. 21
  • Figure 11 compares three features across their maximally activating proteins and visualizes their activations on amsK, the maximally activating protein for feature 9047.Source: Figure 11 · p. 21
  • Feature 6262 is more selective across the three maximally activating glycosyltransferase examples, while the other two features activate on all three.On a separate protein, the latter features activate nearby but distinct regions.Source: E Additional cluster example · pp. 20, 22, Figure 11 · p. 21

F Additional steering experiments

Periodic-glycine features steer glycine predictions in the tested sequences, including alternate variants. Their effects exceed those of random or glycine-specific features but diminish across repeated masked positions.

F Additional steering experiments

  • Each periodic glycine feature steers the model toward a periodic glycine pattern, while combining periodic and non-periodic features has no positive effect on the non-periodic feature.Both feature types increase glycine probability at the steered position in the alternate sequence.Source: Figure 12 · p. 22, F Additional steering experiments · p. 22
  • When applied three at a time, periodic glycine features have a stronger effect on masked-position glycine probability than random or glycine-specific features, with the effect diminishing around the fifth masked position.The comparison was tested across alternate sequence variants and a longer sequence with multiple masked periodic patterns.Source: F Additional steering experiments · p. 22
  • The experiments cover only a few somewhat constrained sequences, leaving the contexts where steering works or fails for further evaluation.Source: F Additional steering experiments · p. 22
Loading 2412.12101v1…