BERTology Meets Biology: Interpreting Attention in Protein Language Models

1 INTRODUCTION . The study of proteins , the fundamental macromolecules governing biology and life itself , has led to remarkable advances in understanding human health and the development of disease therapies . The decreasing cost of sequencing technology has enabled vast databases of naturally occurring proteins ( El-Gebali et al. , 2019a ) , which are rich in information for developing powerful machine learning models of protein sequences . For example , sequence models leveraging principles of co-evolution , whether modeling pairwise or higher-order interactions , have enabled prediction of structure or function ( Rollins et al. , 2019 ) . Proteins , as a sequence of amino acids , can be viewed precisely as a language and therefore modeled using neural architectures developed for natural language . In particular , the Transformer ( Vaswani et al. , 2017 ) , which has revolutionized unsupervised learning for text , shows promise for similar impact on protein sequence modeling . However , the strong performance of the Transformer comes at the cost of interpretability , and this lack of transparency can hide underlying problems such as model bias and spurious correlations ( Niven & Kao , 2019 ; Tan & Celis , 2019 ; Kurita et al. , 2019 ) . In response , much NLP research now focuses on interpreting the Transformer , e.g. , the subspecialty of “ BERTology ” ( Rogers et al. , 2020 ) , which specifically studies the BERT model ( Devlin et al. , 2019 ) . In this work , we adapt and extend this line of interpretability research to protein sequences . We analyze Transformer protein models through the lens of attention , and present a set of interpretability methods that capture the unique functional and structural characteristics of proteins . We also compare the knowledge encoded in attention weights to that captured by hidden-state representations . Finally , we present a visualization of attention contextualized within three-dimensional protein structure . Our analysis reveals that attention captures high-level structural properties of proteins , connecting amino acids that are spatially close in three-dimensional structure , but apart in the underlying sequence ( Figure 1a ) . We also find that attention targets binding sites , a key functional component of proteins ( Figure 1b ) . Further , we show how attention is consistent with a classic measure of similarity between amino acids—the substitution matrix . Finally , we demonstrate that attention captures progressively higher-level representations of structure and function with increasing layer depth . ( a ) Attention in head 12-4 , which targets amino acid pairs that are close in physical space ( see inset subsequence 117D-157I ) but lie apart in the sequence . Example is a de novo designed TIMbarrel ( 5BVL ) with characteristic symmetry . ( b ) Attention in head 7-1 , which targets binding sites , a key functional component of proteins . Example is HIV-1 protease ( 7HVP ) . The primary location receiving attention is 27G , a binding site for protease inhibitor small-molecule drugs . Figure 1 : Examples of how specialized attention heads in a Transformer recover protein structure and function , based solely on language model pre-training . Orange lines depict attention between amino acids ( line width proportional to attention weight ; values below 0.1 hidden ) . Heads were selected based on correlation with ground-truth annotations of contact maps and binding sites . Visualizations based on the NGL Viewer ( Rose et al. , 2018 ; Rose & Hildebrand , 2015 ; Nguyen et al. , 2017 ) . In contrast to NLP , which aims to automate a capability that humans already have—understanding natural language—protein modeling also seeks to shed light on biological processes that are not fully understood . Thus we also discuss how interpretability can aid scientific discovery . 2 BACKGROUND : PROTEINS . In this section we provide background on the biological concepts discussed in later sections . Amino acids . Just as language is composed of words from a shared lexicon , every protein sequence is formed from a vocabulary of amino acids , of which 20 are commonly observed . Amino acids may be denoted by their full name ( e.g. , Proline ) , a 3-letter abbreviation ( Pro ) , or a single-letter code ( P ) . Substitution matrix . While word synonyms are encoded in a thesaurus , proteins that are similar in structure or function are captured in a substitution matrix , which scores pairs of amino acids on how readily they may be substituted for one another while maintaining protein viability . One common substitution matrix is BLOSUM ( Henikoff & Henikoff , 1992 ) , which is derived from co-occurrence statistics of amino acids in aligned protein sequences . Protein structure . Though a protein may be abstracted as a sequence of amino acids , it represents a physical entity with a well-defined three-dimensional structure ( Figure 1 ) . Secondary structure describes the local segments of proteins ; two commonly observed types are the alpha helix and beta sheet . Tertiary structure encompasses the large-scale formations that determine the overall shape and function of the protein . One way to characterize tertiary structure is by a contact map , which describes the pairs of amino acids that are in contact ( within 8 angstroms of one another ) in the folded protein structure but lie apart ( by at least 6 positions ) in the underlying sequence ( Rao et al. , 2019 ) . Binding sites . Proteins may also be characterized by their functional properties . Binding sites are protein regions that bind with other molecules ( proteins , natural ligands , and small-molecule drugs ) to carry out a specific function . For example , the HIV-1 protease is an enzyme responsible for a critical process in replication of HIV ( Brik & Wong , 2003 ) . It has a binding site , shown in Figure 1b , that is a target for drug development to ensure inhibition . Post-translational modifications . After a protein is translated from RNA , it may undergo additional modifications , e.g . phosphorylation , which play a key role in protein structure and function . 3 METHODOLOGY . Model . We demonstrate our interpretability methods on five Transformer models that were pretrained through language modeling of amino acid sequences . We primarily focus on the BERT-Base model from TAPE ( Rao et al. , 2019 ) , which was pretrained on Pfam , a dataset of 31M protein sequences ( ElGebali et al. , 2019b ) . We refer to this model as TapeBert . We also analyze 4 pre-trained Transformer models from ProtTrans ( Elnaggar et al. , 2020 ) : ProtBert and ProtBert-BFD , which are 30-layer , 16-head BERT models ; ProtAlbert , a 12-layer , 64-head ALBERT ( Lan et al. , 2020 ) model ; and ProtXLNet , a 30-layer , 16-head XLNet ( Yang et al. , 2019 ) model . ProtBert-BFD was pretrained on BFD ( Steinegger & Söding , 2018 ) , a dataset of 2.1B protein sequences , while the other ProtTrans models were pretrained on UniRef100 ( Suzek et al. , 2014 ) , which includes 216M protein sequences . A summary of these 5 models is presented in Appendix A.1 . Here we present an overview of BERT , with additional details on all models in Appendix A.2 . BERT inputs a sequence of amino acids x = ( x1 , . . . , xn ) and applies a series of encoders . Each encoder layer ` outputs a sequence of continuous embeddings ( h ( ` ) 1 , . . . , h ( ` ) n ) using a multi-headed attention mechanism . Each attention head in a layer produces a set of attention weights α for an input , where αi , j > 0 is the attention from token i to token j , such that ∑ j αi , j = 1 . Intuitively , attention weights define the influence of every token on the next layer ’ s representation for the current token . We denote a particular head by < layer > - < head_index > , e.g . head 3-7 for the 3rd layer ’ s 7th head . Attention analysis . We analyze how attention aligns with various protein properties . For properties of token pairs , e.g . contact maps , we define an indicator function f ( i , j ) that returns 1 if the property is present in token pair ( i , j ) ( e.g. , if amino acids i and j are in contact ) , and 0 otherwise . We then compute the proportion of high-attention token pairs ( αi , j > θ ) where the property is present , aggregated over a dataset X : pα ( f ) = ∑ x∈X |x|∑ i=1 |x|∑ j=1 f ( i , j ) · 1αi , j > θ /∑ x∈X |x|∑ i=1 |x|∑ j=1 1αi , j > θ ( 1 ) where θ is a threshold to select for high-confidence attention weights . We also present an alternative , continuous version of this metric in Appendix B.1 . For properties of individual tokens , e.g . binding sites , we define f ( i , j ) to return 1 if the property is present in token j ( e.g . if j is a binding site ) . In this case , pα ( f ) equals the proportion of attention that is directed to the property ( e.g . the proportion of attention focused on binding sites ) . When applying these metrics , we include two types of checks to ensure that the results are not due to chance . First , we test that the proportion of attention that aligns with particular properties is significantly higher than the background frequency of these properties , taking into account the Bonferroni correction for multiple hypotheses corresponding to multiple attention heads . Second , we compare the results to a null model , which is an instance of the model with randomly shuffled attention weights . We describe these methods in detail in Appendix B.2 . Probing tasks . We also perform probing tasks on the model , which test the knowledge contained in model representations by using them as inputs to a classifier that predicts a property of interest ( Veldhoen et al. , 2016 ; Conneau et al. , 2018 ; Adi et al. , 2016 ) . The performance of the probing classifier serves as a measure of the knowledge of the property that is encoded in the representation . We run both embedding probes , which assess the knowledge encoded in the output embeddings of each layer , and attention probes ( Reif et al. , 2019 ; Clark et al. , 2019 ) , which measure the knowledge contained in the attention weights for pairwise features . Details are provided in Appendix B.3 . Datasets . For our analyses of amino acids and contact maps , we use a curated dataset from TAPE based on ProteinNet ( AlQuraishi , 2019 ; Fox et al. , 2013 ; Berman et al. , 2000 ; Moult et al. , 2018 ) , which contains amino acid sequences annotated with spatial coordinates ( used for the contact map analysis ) . For the analysis of secondary structure and binding sites we use the Secondary Structure dataset ( Rao et al. , 2019 ; Berman et al. , 2000 ; Moult et al. , 2018 ; Klausen et al. , 2019 ) from TAPE . We employed a taxonomy of secondary structure with three categories : Helix , Strand , and Turn/Bend , with the last two belonging to the higher-level beta sheet category ( Sec . 2 ) . We used this taxonomy to study how the model understood structurally distinct regions of beta sheets . We obtained token-level binding site and protein modification labels from the Protein Data Bank ( Berman et al. , 2000 ) . For analyzing attention , we used a random subset of 5000 sequences from the training split of the respective datasets ( note that none of the aforementioned annotations were used in model training ) . For the diagnostic classifier , we used the respective training splits for training and the validation splits for evaluation . See Appendix B.4 for additional details . Experimental details We exclude attention to the [ SEP ] delimiter token , as it has been shown to be a “ no-op ” attention token ( Clark et al. , 2019 ) , as well as attention to the [ CLS ] token , which is not explicitly used in language modeling . We only include results for attention heads where at least 100 high-confidence attention arcs are available for analysis . We set the attention threshold θ to 0.3 to select for high-confidence attention while retaining sufficient data for analysis . We truncate all protein sequences to a length of 512 to reduce memory requirements.1 We note that all of the above analyses are purely associative and do not attempt to establish a causal link between attention and model behavior ( Vig et al. , 2020 ; Grimsley et al. , 2020 ) , nor to explain model predictions ( Jain & Wallace , 2019 ; Wiegreffe & Pinter , 2019 ) .

Summary: The authors analyzed how attention and embeddings of Transformers trained on protein sequence correlate with protein properties such as pairwise contacts, binding sites, and post-translational modifications. The paper extends existing papers such as Rives 2020 ‘Biological structure and function emerge…’ by showing that layers learn more and more complex protein features with increasing layer depth and by proposing new visualization techniques. The paper is mostly clearly written while the methodological contributions are incremental. The evaluation needs to be strengthened.