TLDR: ProtSAE is a new method that improves the interpretability of protein language models (PLMs) by using semantically-guided sparse autoencoders (SAEs). It tackles the problem of “semantic entanglement” in traditional SAEs by incorporating both annotation data and protein domain knowledge during training. This allows ProtSAE to learn more biologically relevant and disentangled features, leading to better performance in protein function prediction and enabling more effective steering of PLMs for generating proteins with desired characteristics.
Protein Language Models (PLMs) have rapidly advanced in recent years, finding applications in crucial areas like protein function prediction, structural modeling, and protein design. However, understanding their internal workings remains a significant challenge. This lack of interpretability is a major hurdle, especially for protein engineering, where knowing how latent features map to biological concepts (like binding pockets or fold families) is essential for improving model performance, robustness, and gaining new biological insights.
Sparse Autoencoders (SAEs) have emerged as a promising tool for interpreting large language models, including PLMs. SAEs work by decomposing the hidden representations of these models into sparse features, which can then be analyzed for their correlation with specific biological concepts. This allows for a degree of interpretability and even the ability to steer model generation by selectively activating these features.
However, a key limitation of traditional SAEs is “semantic entanglement.” This means that individual neurons within the SAE often mix multiple biological concepts, making it difficult to reliably interpret what a specific neuron represents or how to precisely manipulate model behaviors. Imagine a single neuron trying to represent both “binding pocket” and “post-translational modification” simultaneously – its meaning becomes ambiguous.
To address this challenge, researchers have introduced ProtSAE, a novel semantically-guided Sparse Autoencoder. Unlike existing SAEs that rely on post-training annotation datasets to filter and interpret activations, ProtSAE integrates semantic disentanglement directly into its training process. It achieves this by leveraging both annotation datasets and rich protein domain knowledge.
ProtSAE’s approach involves two main guiding principles. First, it uses semantic annotations to constrain the relationship between defined activations and specific concepts during training. This ensures that certain neurons are selectively activated only by proteins associated with a particular concept. To make sure these concept-specific activations actively contribute to the model’s output, ProtSAE employs “forced activations” and “feature rescaling.” This prevents the model from relying solely on entangled, unsupervised activations.
Second, recognizing that biological concepts are often interconnected, ProtSAE incorporates protein domain knowledge using a method called ELEmbeddings. This allows the model to learn logical constraints among concepts, such as “is-a” or “part-of” relationships, enhancing the interpretability and semantic consistency of the learned features. The overall training objective for ProtSAE combines reconstruction accuracy, supervised prediction based on annotations, and semantic regularization guided by this domain knowledge.
Extensive experiments demonstrate ProtSAE’s effectiveness. In interpretability visualizations, ProtSAE successfully highlights biologically relevant features, such as iron ion binding regions on TonB-dependent receptors, sodium ion transport on transmembrane segments, and specific metal ion binding sites. This visual evidence shows a strong alignment between the learned features and actual protein structures and functions.
Quantitative evaluations further support these findings. ProtSAE consistently achieves higher F1-scores in relevance-based interpretation, indicating that its features are more semantically aligned with target concepts compared to baseline SAEs. In probing-based interpretation tasks, such as protein function prediction, ProtSAE outperforms all SAE baselines and even a dictionary learning method called SpLiCE, achieving performance comparable to direct linear probing on the Protein Language Model’s hidden representations. This suggests that ProtSAE effectively mitigates semantic loss during training.
Performance analyses across varying levels of sparsity show that ProtSAE maintains superior predictive performance while preserving high reconstruction fidelity. An ablation study confirmed the importance of each component: detaching prediction weights, incorporating axiom-based domain knowledge, using forced activations, and feature rescaling all contribute significantly to ProtSAE’s success in disentangling semantics and improving interpretability.
Beyond interpretation, ProtSAE also demonstrates potential in steering PLMs for downstream generation tasks. By selectively activating concept-specific features, ProtSAE can guide the generation of new protein sequences. Steering experiments show that generated proteins exhibit significantly improved structural similarity to natural proteins with the target concepts, as measured by TM-score and RMSD, and also show increased structural stability (pLDDT scores). For instance, ProtSAE successfully generated a protein structurally similar to a natural DNA-binding transcription repressor, while maintaining sequence novelty. You can find more details about this research in the full paper available here.
Also Read:
- Unveiling the Inner Workings: How AI Models Fabricate Information from Noise
- Bridging Language Models and Knowledge Graphs for Biomedical Understanding
In conclusion, ProtSAE offers a significant advancement in interpreting and controlling protein language models. By addressing semantic entanglement through guided training with both annotation data and domain knowledge, it yields more biologically relevant and interpretable hidden features, paving the way for more precise protein engineering and deeper biological insights.
