Identifying Functionally Important Residues: BERT Entropy and Interpretability Applied to Protein Function Prediction

Introduction: Predicting protein function from sequence is a longstanding challenge in bioinformatics. Various models have been developed to predict Gene Ontology (GO) terms, which describe the molecular functions, biological processes, and cellular components of proteins. GO provides a structured vocabulary to annotate proteins with functional information. In this project, we aim to interpret a GO prediction model by identifying residues crucial for predicting specific GO functions. Traditionally, sequence conservation, measured through Shannon entropy in multiple sequence alignment (MSA), has been used to infer residue importance. However, here we introduce an innovative approach by incorporating BERT —a model trained to demask hidden residues—to estimate the conditional distribution of amino acids for a given sequence. This provides insights into residue importance, with lower entropy indicating higher conservation and potential functional significance. We mutated the important residues identified through this to use Local Interpretable Model-agnostic Explanations (LIME) to reproduce a strong machine learning model with a limited scope, but more interpretable, logistic regression classifier.

Materials and

Methods: We fine-tuned ESM-2, a language model for protein-related tasks, for GO classification with 650 million parameters. The Mechanism and Catalytic Site Atlas (M-CSA) provided data for identifying experimentally validated catalytic residues, resulting in a dataset of 785 annotated proteins with associated GO terms. We ran the following with each protein from this dataset. To estimate the conditional distribution of amino acids, we tokenized the target protein and generated randomly masked sequences by masking 15% of the sequence each time. The masked sequences are then used to train ESM-2. Averaging the distributions for each residue, we calculated the entropy of residues per protein, and evaluated how potentially high conservation residues from this compare to annotated ones by using mean reciprocal rank (MRR). The top 75 most important residues determined from this were mutated in performing LIME. LIME generates local explanations for model predictions by approximating their behavior through simpler, interpretable models. For each protein, we generated several hundred mutated sequences and passed them all through the GO prediction model. Then, we embedded the proteins with four chemical properties—aromaticity, pKa, hydrophobicity, and molecular weight—before applying logistic regression. The weights obtained from the logistic regression revealed the importance of each residue in outputting the GO classification for that protein. Finally, we compared this with annotated residues from the M-CSA dataset using MRR.

Results, Conclusions, and Discussions: We evaluate residue importance as a retrieval problem, in which the goal is to assign the most importance to residues relevant to a particular function. Retrieving based on BERT entropy alone results in a mean-reciprocal rank (MRR) of 0.14, indicating that one needs to go through 7 residues on average to find one that is important. LIME performed worse with a MRR of 0.062. However, LIME gives significantly greater weights to the annotated residues when run on its correct GO term versus random GO terms, reflecting that LIME actually indicates important sites for functions they are actually responsible for. BERT entropy offers a novel perspective on residue importance. Used in conjunction with LIME, we can determine which sites on a protein most likely correspond to certain functions. This work allows us to potentially, given a protein, on top of predicting the GO class(es) from a black box GO prediction model, we can also identify specific residues which are responsible for this function(s). We could also use this to estimate the effect in function given a certain mutation in a protein.

Acknowledgements (Optional):