From Mechanistic Interpretability to Mechanistic Biology: Training, Evaluating, and Interpreting Sparse Autoencoders on Protein Language Models

We thank the reviewer for their suggestions and enthusiasm for our work.

We agree that one of the limitations of the paper is a heavy reliance on qualitative measures of feature interpretability. Towards making feature analysis more quantitative, we introduced the family specificity and activation pattern categorizations. In our experience, family specific-features are interpretable (can be interpreted as a specific sequence motif in a specific family), so the number of seemingly interpretable features is lower bounded by the number of family-specific features. But these measures admittedly don’t truly measure interpretability, and we would love to include results for human raters, and compare it to the ESM baseline. As done in the Gated SAE paper, we will conduct a blinded human rater experiment where our 5 raters who are familiar with protein biology will assess the interpretability of an SAE latent and ESM baseline as being interpretable (yes/maybe/no). We will include the results in a revision.

Your interpretation of our goal in the probing section is correct; we are aiming to show that SAE representations haven’t lost key information (as demonstrated by task performance) and that the highest weighted latents correspond to features which makes biological sense with respect to the task. You are right that we could identify relevant features using a simpler method. Our motivation for probing, besides simply identifying relevant features for a task, was to demonstrate how SAE embeddings could be used as a drop-in replacement for dense ESM embeddings, without a significant hit to performance. Since in the pLM literature, linear probing is typically used to measure downstream task performance [1], we opted to do the same to make a head-on comparison.

Regarding steering family specific latents, you make an excellent point. We will add a control where we steer a different (not family specific) latent to better demonstrate that family specific features are more important for conserved regions than other features. We will include this in the camera ready version. This is our attempt at getting more “causal” evidence to interpret SAE latents. As touched upon in our future directions section, getting interpretable steering results using SAE latents beyond simple amino acid specific features has remained challenging, and is a direction we will continue to explore in the future.

In our calculation of the F1 score for family specificity, we also include sequences which do not activate the latent of interest. We believe this should help account for the false negative problem you raise.

To make our work more accessible to the broader mechanistic interpretability, we will add a primer to introduce protein language models and some domain specific terms we often use in the paper. There have been some good reviews recently published which we can link as well [2]. And we’ve happy you like the name of our project!

[1] https://www.biorxiv.org/content/10.1101/2024.02.05.578959v2 [2] https://www.nature.com/articles/s41587-024-02123-4

We thank the reviewer for their detailed reading and thoughts.

We acknowledge that the lack of other pLMs in our analysis beyond ESM-2 means that we should not make general conclusions about all pLMs. We intend to update the text to reflect this, replacing pLM with ESM-2 to narrow the scope of the results.

As the reviewer points out, the performance of SAE probes on the mammalian expression task lags behind the ESM probes. In the text, we say that “linear probes on SAEs achieve performance similar to their ESM embedding baselines across all layers”. We intend to update the text to qualify this statement, that they achieve similar performance on most tasks.

In the related work, the authors argue that their work focuses on coevolution and scientific discovery. I don't see any experiments focusing on coevolution; am I misunderstanding?

This is a good point, we could have been more specific. By coevolution, we refer to our analysis of features which seem to correspond to protein family-specific motifs.

We would like to expand on why we view our results as significant from a biological perspective. Understanding biological sequences such as protein sequences remains a difficult problem. Traditionally, biologists have used tools such as multiple sequence alignments and statistical models to identify local patterns in protein sequences. Recently, protein language models trained on millions of sequences across evolution have been successful in providing useful protein representations. Embeddings from pLMs are often used as input to simple models (e.g. linear regression) on downstream tasks such as protein structure and variant effect prediction. Yet as these representations are uninterpretable, it is not easy to identify what features the pLM has learned which contribute to task performance. We advance on prior work by demonstrating SAEs can help reveal what ESM-2 learned features are important for downstream task performance. For example, by probing on SAE latents, we can identify that amino acid composition is an important predictor for thermostability. This may allow a biologist to gain insight into novel sequence determinants of a protein property which would have required manual feature crafting previously. As representations from ESM are widely used, we offer SAE representations as a drop-in replacement with (usually) similar performance as the ESM representations but with interpretable dimensions. We hope this clarifies the advantage of SAE-derived representations.

We thank the reviewer for these thoughtful comments and suggestions.

We created an anonymous link for our InterProt visualizer at http://icml.interprot.com and hope that it can provide context on our manual interpretation process and showcase some interpretable features. To support our claims around interpretability, we agree with the need for more human label data and quantitative analysis. As done in the Gated SAE paper (Rajamanoharan et al.), we will conduct a blinded human rater experiment where 5 raters who are familiar with protein biology will assess the interpretability of an SAE latent and ESM baseline as being interpretable (yes/maybe/no). We will include the results in a revision. We hope that this can provide more robust evidence on the value added by SAE compared to directly analyzing the base model.

We agree with the lack of details on how our work differs from Simon et al. and plan to include more information in our revision. The key differences are as follows:

• We trained TopK SAEs while Simon et al. used ReLU SAEs.
• We used a larger, 650M param variant ESM-2, compared to the 8M variant used by Simon et al. The 650M param model is far more widely used.
• We propose a framework for systematically categorizing latents by family-specificity and activation patterns; Simon et al. does not do latent categorization. We also evaluate the effect of different SAE hyperparameters on these feature categorizations. We use linear probes on four downstream tasks to extract interpretable features with the goal of enabling scientific discovery. Simon et al. focuses on demonstrating the feasibility of SAEs on pLMs and performs quantitative interpretability comparisons between ESM and SAE.
• We share an open-source visualizer InterProt with data visualizations aimed for manual feature interpretation. For example, each SAE latent is displayed with a collection activating sequences across different activating ranges, whether they cluster within specific protein families, providing options to align them. InterProt also enables searching a sequence across all latents, a feature that has enabled discovery of interesting, interpretable latents starting from a protein of interest. Simon et al. also provides a visualizer though it focuses more on displaying activating distributions of each latent and whether it has been linked to any Swiss-Prot concepts. Simon et al. proposed a method to automate the interpretation of SAE features using an LLM. We did not explore this approach.

We see models like ProtoPNet with built-in interpretability as an exciting and complementary direction to post-hoc interpretation (of usually larger and more performant pLMs) via SAEs. Concept bottleneck models have also been applied to proteins, and have shown competitive performance with pLMs [2]. While we agree that a comparison of SAEs to these methods is outside the scope of this work, we plan to add a discussion of these references to our Related Works section.

Finally, we thank the reviewer for the reference on the previous work by Makelov et al. on the effects of training data distribution, a direction we plan to explore in follow-up work.

[1] https://arxiv.org/abs/2404.16014 [2] https://arxiv.org/abs/2411.06090

We thank the reviewer for the thorough review and many good suggestions.

We agree with the reviewer’s points around the limitations of SAEs. The presence of a large number of family-specific features suggests that SAEs do indeed learn, or memorize, MSAs. Furthermore, the activation pattern of family-specific features support the hypothesis proposed by Zhang et al. [1] that pLMs store evolutionary statistics via common motifs. Although one cannot conclude from the analysis of SAE features that pLMs have learned any biophysics (as Zhang et al. suggest, they likely do not), the presence of non-family-specific features do present evidence that ESM2 has learned some patterns beyond memorizing common motifs within protein families. For example, features corresponding to generic helixes (L16/2611), hydrophobic core residues (L8/2064), and nuclear localization signals (L28/2375) suggest that pLMs do contain generic notions not specific to a protein family. Our mental model is that pLMs not only store co-evolutionarily information from MSAs but also some useful statistical interpolations between them.

As is common in interpretable ML, understanding SAE latent representations still relies heavily on human intuition. In the case of pLM SAEs, we benefit from established biological knowledge—such as secondary structure and known motifs—which gives us useful reference points. It's encouraging that many SAE latents align with these expected features. However, our claims would be stronger with a quantitative comparison showing that SAE latents are more interpretable than an ESM baseline. To that end:

• Though using a smaller model in the ESM family and with a different architecture, Simon et al. [2] showed that more SAE latents map to SwissProt annotations compared to ESM neurons. For a more comprehensive comparison, see our response to reviewer 5Na5.
• In A.7., we show some analysis of linear probe weight distributions that weakly suggest SAE probe coefficients to be more interpretable.
• As done in the Gated SAE paper, we will conduct a blinded human rater experiment where 5 raters who are familiar with protein biology will assess the interpretability of an SAE latent and ESM baseline as being interpretable (yes/maybe/no). We will include the results in a revision.

Recent criticism of SAEspoint to the limitations of LLM-based auto-interpretation techniques [3], which were not used in our work. It has also been found that SAEs trained on the same data learn different features [15].

Many other mechanistic interpretability techniques have been applied to pLMs. For example, attention matrices have been shown to resemble pairwise contact maps ([4], [5], [6]), and contain information about motifs such as signal peptides [7], which are also represented as SAE features. More recent methods such as model diffing [8], trans-coders/cross-coders ([9], [10]), and other SAE architectures ([11], [12]) are promising future directions for this work.

We agree that mean-pooling SAE embeddings when performing downstream probes has important limitations. For the presence of a short functional motif or residue-level signal, max-pooling will likely yield clearer results. We would also like to try more sophisticated pooling methods such as aggregating via optimal transport [14] in the future.

For the definition of family-specific features, we set the F1 threshold to 0.7 because we observed there to be a significant dropoff in F1 score below that point. We will update the supplement with a figure demonstrating this.

Finally, we agree with the suggestion to include more quantitative evaluations of SAE features. Doing so across different hyperparameter schemes would additionally enable a more principled hyperparameters selection strategy. Our initial hope was to achieve this via our feature categorization analysis but only observed small effects across different hyperparameters (A.4). We will evaluate our probes across SAEs trained with different hyperparameters. Specifically, a range of SAE latent hidden dimensions and a range of k values. We hope this will help demonstrate the robustness of our results.

[1] https://tinyurl.com/m8a5ratb [2] https://tinyurl.com/3xsfefaj
[3] https://arxiv.org/abs/2501.17727 [4] https://arxiv.org/abs/2404.16014 [5] https://arxiv.org/abs/2006.15222 [6] https://tinyurl.com/yja6xyhf [7] https://tinyurl.com/mub4tuy5 [8] https://tinyurl.com/4hwm4tp8
[9] https://arxiv.org/abs/2406.11944 [10] https://tinyurl.com/4hwm4tp8 [11] https://arxiv.org/abs/2404.16014 [12] https://arxiv.org/abs/2407.14435 [13] https://tinyurl.com/yc2memtc [14] https://tinyurl.com/yc8j3fzt [15] https://arxiv.org/abs/2501.16615