Interpreting and Steering Protein Language Models through Sparse Autoencoders

2025-02-13

AI Summary

Understanding the internal representational mechanisms of protein language models (PLMs) remains challenging. Method: We introduce the first sparse autoencoder (SAE)-based interpretability framework tailored for PLMs—specifically ESM-2 (8M)—that maps hidden-layer activations to biologically annotated functional motifs (e.g., transmembrane regions, binding sites, zinc finger domains). Statistical significance of function-specific neurons is assessed via Fisher’s exact test. Furthermore, we develop a conditional generation strategy guided by activations of identified interpretable neurons to enable targeted sequence design of target motifs (e.g., signal peptides, zinc fingers). Contribution/Results: Our approach successfully identifies multiple biologically grounded, interpretable neurons and achieves over threefold enrichment of target motifs in generated sequences. This work establishes a novel paradigm for mechanistic understanding and controllable design using PLMs.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization

Abstract

The rapid advancements in transformer-based language models have revolutionized natural language processing, yet understanding the internal mechanisms of these models remains a significant challenge. This paper explores the application of sparse autoencoders (SAE) to interpret the internal representations of protein language models, specifically focusing on the ESM-2 8M parameter model. By performing a statistical analysis on each latent component's relevance to distinct protein annotations, we identify potential interpretations linked to various protein characteristics, including transmembrane regions, binding sites, and specialized motifs. We then leverage these insights to guide sequence generation, shortlisting the relevant latent components that can steer the model towards desired targets such as zinc finger domains. This work contributes to the emerging field of mechanistic interpretability in biological sequence models, offering new perspectives on model steering for sequence design.

Problem

• Interpreting protein language models
• Steering sequence generation
• Mechanistic interpretability in biology

Innovation

• Sparse autoencoders interpret protein models
• Statistical analysis links latent to proteins
• Latent components guide sequence generation