ProtSAE: Disentangling and Interpreting Protein Language Models via Semantically-Guided Sparse Autoencoders

TL;DR

ProtSAE tackles semantic entanglement in sparse autoencoders for protein language models by introducing semantically guided training that leverages annotation data and domain knowledge through ELEmbeddings. The method defines concept-specific activations, ties encoder directions to ontology-informed embeddings, and enforces semantic constraints via axiom losses while preserving reconstruction fidelity. Empirical results show improved interpretability, competitive predictive performance across GO ontologies and metal-binding tasks, and the ability to steer PLM generation toward concept-aligned structures. This approach enhances mechanistic understanding and controllability of PLMs in protein biology, enabling more reliable interpretation and targeted protein design.

Abstract

Sparse Autoencoder (SAE) has emerged as a powerful tool for mechanistic interpretability of large language models. Recent works apply SAE to protein language models (PLMs), aiming to extract and analyze biologically meaningful features from their latent spaces. However, SAE suffers from semantic entanglement, where individual neurons often mix multiple nonlinear concepts, making it difficult to reliably interpret or manipulate model behaviors. In this paper, we propose a semantically-guided SAE, called ProtSAE. Unlike existing SAE which requires annotation datasets to filter and interpret activations, we guide semantic disentanglement during training using both annotation datasets and domain knowledge to mitigate the effects of entangled attributes. We design interpretability experiments showing that ProtSAE learns more biologically relevant and interpretable hidden features compared to previous methods. Performance analyses further demonstrate that ProtSAE maintains high reconstruction fidelity while achieving better results in interpretable probing. We also show the potential of ProtSAE in steering PLMs for downstream generation tasks.

Figures (14)

• Figure 1: Illustration of SAE semantic entanglement (left): individual neurons conflate multiple biological concepts, and semantic disentanglement (right): each defined neuron maps to a single biological concept.
• Figure 2: An overview of ProtSAE (left) and the baseline SAE (right). In the baseline SAE, annotation data is used post hoc to interpret learned features via relevance-based and probing-based methods. In contrast, ProtSAE incorporates semantic guidance during training by leveraging annotation data and protein domain knowledge to achieve semantic disentanglement. It also uses forced activations and feature rescaling to learn meaningful features while preserving reconstruction fidelity.
• Figure 3: Interpretability visualization shows that ProtSAE reveals semantic alignment between learned features and protein structures, including functional regions and ion binding sites. We use red intensity to indicate feature activation strength, and green sticks to mark ground truth binding sites.
• Figure 4: Comparison of relevance-based interpretation
• Figure 5: Performance comparison under different sparsity on the BPO dataset