From Mechanistic Interpretability to Mechanistic Biology: Training, Evaluating, and Interpreting Sparse Autoencoders...

Area of Science:

• Computational Biology
• Bioinformatics
• Machine Learning in Biology

Background:

• Protein language models (pLMs) are advanced tools for predicting protein structure and function, trained on vast sequence datasets.
• The internal features or representations learned by pLMs are not fully understood, limiting insights into their predictive power and potential for biological discovery.
• Understanding pLM features could illuminate how these models capture biological information and potentially uncover new protein biology.

Purpose of the Study:

• To investigate the specific features learned by protein language models (pLMs).
• To characterize the nature of these features (generic vs. specific) and their relationship to protein properties.
• To explore the utility of these features for hypothesis generation regarding unknown biological mechanisms.

Main Methods:

• Training sparse autoencoders (SAEs) on the residual stream outputs of the ESM-2 pLM.
• Characterizing the learned SAE features to understand their representation of protein sequence information.
• Employing linear probing on SAE features to identify sequence determinants of protein properties like thermostability and localization.
• Developing visualization tools for interpreting predictive SAE features.

Main Results:

• Demonstrated that pLMs utilize a mix of generic and family-specific features to represent proteins.
• Successfully identified known sequence determinants for protein thermostability and subcellular localization using linear probing of SAE features.
• Highlighted predictive features lacking clear functional associations, suggesting potential roles in undiscovered biological mechanisms.

Conclusions:

• pLM features offer a blend of general and specific biological information.
• SAE analysis provides a method to interpret pLM representations and link them to protein functions.
• This approach enhances understanding of pLM limitations and facilitates the generation of new biological hypotheses.