Original: BERT OLOGY M EETS B IOLOGY : I NTERPRETING A TTENTION IN P ROTEIN L ANGUAGE M ODELS
Revised: BERT OLOGY M EETS B IOLOGY : I NTERPRETING A TTENTION IN P ROTEIN L ANGUAGE M ODELS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Transformer architectures have proven to learn useful representations for protein classiﬁcation and generation tasks.
Revised: Transformer architectures have proven to learn useful representations for protein classiﬁcation and generation tasks.

Original: However, these representations present challenges in interpretability.
Revised: However, these representations present challenges in interpretability.

Original: In this work, we demonstrate a set of methods for analyzing protein transformer models through the lens of attention.
Revised: In this work, we demonstrate a set of methods for analyzing protein Transformer models through the lens of attention.

Original: We show that attention: (1) captures the folding structure of proteins, connecting amino acids that are far apart in the underlying sequence, but spatially close in the three-dimensional structure, (2) targets binding sites, a key functional component of proteins and (3) focuses on progressively more complex biophysical properties with increasing layer depth.
Revised: We show that attention: (1) captures the folding structure of proteins, connecting amino acids that are far apart in the underlying sequence, but spatially close in the three-dimensional structure, (2) targets binding sites, a key functional component of proteins and (3) focuses on progressively more complex biophysical properties with increasing layer depth.

Original: We ﬁnd this behavior to be consistent across three Transformer architectures (BERT, ALBERT, XLNet) and two distinct protein datasets.
Revised: We ﬁnd this behavior to be consistent across three Transformer architectures (BERT, ALBERT, XLNet) and two distinct protein datasets.

Original: We also present a three-dimensional visualization of the interaction between attention and protein structure.
Revised: We also present a three-dimensional visualization of the interaction between attention and protein structure.

Original: Code for visualization and analysis is available at [supplementary-material] .
Revised: Code for visualization and analysis is available at [supplementary-material] .

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: Proteins, as a sequence of amino acids, can be viewed precisely as a language and therefore modeled using neural architectures developed for natural language.
Revised: Proteins, as a sequence of amino acids, can be viewed precisely as a language and therefore modeled using neural architectures developed for natural language.

Original: In particular, the Transformer (Vaswani et al., 2017), which has revolutionized unsupervised learning for text, shows promise for similar impact on protein sequence modeling.
Revised: In particular, the Transformer (Vaswani et al., 2017), which has revolutionized unsupervised learning for text, shows promise for similar impact on protein sequence modeling.

Original: 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).
Revised: 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).

Original: In response, much NLP research now focuses on interpreting the Transformer, e.g., the subspecialty of “BERTology” (Rogers et al., 2020), which speciﬁcally studies the BERT model (Devlin et al., 2019).
Revised: In response, much NLP research now focuses on interpreting the Transformer, e.g., the subspecialty of “BERTology” (Rogers et al., 2020), which speciﬁcally studies the BERT model (Devlin et al., 2019).

Original: In this work, we adapt and extend this line of interpretability research to protein sequences.
Revised: In this work, we adapt and extend this line of interpretability research to protein sequences.

Original: 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.
Revised: 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.

Original: We also introduce two analysis methods that apply to Transformer models more broadly: a novel calibration analysis of attention as a measure of probability, and a joint, cross-layer probing analysis of attention weights and embeddings, showing that layers of the model in which knowledge is accrued in embeddings may be far removed from layers in which this information is leveraged in the attention mechanism.
Revised: We also perform a joint, cross-layer probing analysis of attention weights and embeddings, showing that layers of the model in which knowledge is accrued in embeddings may be far removed from layers in which this information is leveraged in the attention mechanism.

Original: Finally, we present a visualization of attention contextualized within three-dimensional protein structure.
Revised: Finally, we present a visualization of attention contextualized within three-dimensional protein structure.

Original: 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).
Revised: 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).

Original: We also ﬁnd that attention targets binding sites, a key functional component of proteins (Figure 1b).
Revised: We also ﬁnd that attention targets binding sites, a key functional component of proteins (Figure 1b).

Original: Further, we show that attention captures substitution properties of  amino acids, and constructs progressively higher-level representations of structure and function with increasing 
Revised: 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.

Original: [Figure]
Revised: [Figure]

Original: layer depth. 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.
Revised: 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.

Original: Thus we also discuss how interpretability can aid scientiﬁc discovery.
Revised: Thus we also discuss how interpretability can aid scientiﬁc discovery.

Original:  B ACKGROUND : P ROTEINS
Revised:  B ACKGROUND : P ROTEINS

Original: In this section we provide background on the biological concepts discussed in later sections.
Revised: In this section we provide background on the biological concepts discussed in later sections.

Original: Amino acids.
Revised: Amino acids.

Original: Just as language is composed of words from a shared lexicon, every protein sequence is formed from a vocabulary of20 standard amino acids.
Revised: 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.

Revised: Amino acids may be denoted by their full name, e.g. Proline , a 3 letter abbreviation, e.g. Pro , or a single-letter code, e.g. P .

Original:  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.
Revised: 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.

Original: One common substitution matrix is BLOSUM (Henikoff & Henikoff, 1992), which is derived from co-occurrence statistics of amino acids in aligned protein sequences.
Revised: One common substitution matrix is BLOSUM (Henikoff & Henikoff, 1992), which is derived from co-occurrence statistics of amino acids in aligned protein sequences.

Original: Protein structure.
Revised: Protein structure.

Original: Though a protein may be abstracted as a sequence of amino acids, it represents a physical entity with a well-deﬁned three-dimensional structure (Figure 1).
Revised: Though a protein may be abstracted as a sequence of amino acids, it represents a physical entity with a well-deﬁned three-dimensional structure (Figure 1).

Original: Secondary structure describes the local segments of proteins, and may be divided into three broad categories: Helix ,
Revised: Secondary structure describes the local segments of proteins; two commonly observed types are the alpha helix and beta sheet .

Original: Strand , and Turn/Fold , as well as an Other category for local structure that falls outside these categories.

Original: Tertiary structure encompasses the large-scale formations that determine the overall shape and function of the protein.
Revised: Tertiary structure encompasses the large-scale formations that determine the overall shape and function of the protein.

Original: 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).
Revised: 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).

Original: Binding sites.
Revised: Binding sites.

Original: Proteins may also be characterized by their functional properties.
Revised: Proteins may also be characterized by their functional properties.

Original: Binding sites are protein regions that bind with other molecules  to carry out a speciﬁc function.
Revised: Binding sites are protein regions that bind with other molecules (proteins, natural ligands, and small-molecule drugs) to carry out a speciﬁc function.

Original: For example, the HIVprotease is an enzyme responsible for a critical process in replication of HIV (Brik & Wong, 2003).
Revised: For example, the HIV-1 protease is an enzyme responsible for a critical process in replication of HIV (Brik & Wong, 2003).

Original: It has a binding site, shown in Figure 1b, that is a target for drug development to ensure inhibition.
Revised: It has a binding site, shown in Figure 1b, that is a target for drug development to ensure inhibition.

Original: Post-translational modiﬁcations.
Revised: Post-translational modiﬁcations.

Original: After a protein is translated from RNA, it may undergo additional modiﬁcations, e.g. phosphorylation, which play a key role in protein structure and function.
Revised: After a protein is translated from RNA, it may undergo additional modiﬁcations, e.g. phosphorylation, which play a key role in protein structure and function.

Original:  M ETHODOLOGY
Revised:  M ETHODOLOGY

Original: Model.
Revised: Model.

Original: We demonstrate our interpretability methods on ﬁve Transformer models that were pretrained through language modeling of amino acid sequences.
Revised: We demonstrate our interpretability methods on ﬁve Transformer models that were pretrained through language modeling of amino acid sequences.

Original: 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).
Revised: 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 (El-Gebali et al., 2019b).

Original: We refer to this model as TapeBert .
Revised: We refer to this model as TapeBert .

Original: 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., 2019) model, and
Revised: 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.) model, and

Original: ProtXLNet , a 30-layer, 16-head XLNet (Yang et al., 2019) model.
Revised: ProtXLNet , a 30-layer, 16-head XLNet (Yang et al., 2019) model.

Original: ProtBert-BFD was pretrained on
Revised: ProtBert-BFD was pretrained on

Original: 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.
Revised: 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.

Original: Here we present an overview of BERT, with additional details on all models in Appendix A.1.
Revised: Here we present an overview of BERT, with additional details on all models in Appendix A.1.

Original: BERT inputs a sequence of amino acids x = ( x 1 , . . .
Revised: BERT inputs a sequence of amino acids x = ( x 1 , . . .

Original: , x n ) and applies a series of encoders.
Revised: , x n ) and applies a series of encoders.

Original: Each encoder layer (cid:96) outputs a sequence of continuous embeddings ( h ( (cid:96) ) 1 , . .
Revised: Each encoder layer (cid:96) outputs a sequence of continuous embeddings ( h ( (cid:96) ) 1 , . .

Original: , h ( (cid:96) ) n ) using a multi-headed attention mechanism.
Revised: , h ( (cid:96) ) n ) using a multi-headed attention mechanism.

Original: 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 (cid:80)
Revised: 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 (cid:80)

Original: j α i,j = 1 .
Revised: j α i,j = 1 .

Original: Intuitively, attention weights deﬁne the inﬂuence of every token on the next layer’s representation for the current token.
Revised: Intuitively, attention weights deﬁne the inﬂuence of every token on the next layer’s representation for the current token.

Original: We denote a particular head by <layer>-<head_index> , e.g. head 3-7 for the 3rd layer’s 7th head.
Revised: We denote a particular head by <layer>-<head_index> , e.g. head 3-7 for the 3rd layer’s 7th head.

Original: Attention analysis.
Revised: Attention analysis.

Original: We analyze how attention aligns with various protein properties.
Revised: We analyze how attention aligns with various protein properties.

Original: For properties of token pairs, e.g. contact maps, we deﬁne an indicator function f
Revised: For properties of token pairs, e.g. contact maps, we deﬁne an indicator function f

Original: ( 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.
Revised: ( 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.

Original: We then compute the proportion of high-attention token pairs ( α i,j > θ ) where the property is present, aggregated over a dataset X :
Revised: We then compute the proportion of high-attention token pairs ( α i,j > θ ) where the property is present, aggregated over a dataset X :

Original: [Equation]
Revised: [Equation]

Original: where θ is a threshold to select for high-conﬁdence attention weights.
Revised: where θ is a threshold to select for high-conﬁdence attention weights.

Original:  For properties of individual tokens, e.g. binding sites, we substitute f ( j ) for f ( i, j ) in Eq. 1, where f ( j ) returns 1 if the property is present in token j (e.g. if j is a binding site).
Revised: We also present an alternative, continuous version of this metric in Appendix B.1. For properties of individual tokens , e.g. binding sites, we deﬁne f ( i, j ) to return 1 if the property is present in token j (e.g. if j is a binding site).

Original: In this case, p α ( f ) equals the proportion of attention that is directed to  property f .
Revised: 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).

Revised: When applying these metrics, we include two types of checks to ensure that the results are not due to chance.

Original: We present an alternative, continuous version of this metric in Appendix B.1.
Revised: First, we test that the proportion of attention that aligns with particular properties is signiﬁcantly 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 shufﬂed attention weights. We describe these methods in detail in Appendix B.2.

Original: Probing tasks.
Revised: Probing tasks.

Original: We also perform probing tasks on the model, which test the knowledge contained in model representations by using the them as inputs to a classiﬁer that predicts a property of interest (Veldhoen et al., 2016;
Revised: We also perform probing tasks on the model, which test the knowledge contained in model representations by using the them as inputs to a classiﬁer that predicts a property of interest (Veldhoen et al., 2016;

Original: Conneau et al., 2018; Adi et al., 2016).
Revised: Conneau et al., 2018; Adi et al., 2016).

Original: The performance of the probing classiﬁer serves as a measure of the knowledge of the property that is encoded in the representation.
Revised: The performance of the probing classiﬁer serves as a measure of the knowledge of the property that is encoded in the representation.

Original: We run both embedding probes , which assess the knowledge encoded in the output embeddings of each layer, and attention probes (Coenen et al., 2019; Clark et al., 2019), which measure the knowledge contained in the attention weights for pairwise features.
Revised: 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.

Original: Details are provided in Appendix B.2.
Revised: Details are provided in Appendix B.3.

Original: Datasets.
Revised: Datasets.

Original: 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).
Revised: 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).

Original: 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. 
Revised: 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 considered a more ﬁne-grained 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 in to study how the model understood structurally distinct regions of beta sheets.

Original: We obtained token-level binding site and protein modiﬁcation labels from the Protein Data Bank (Berman et al., 2000).
Revised: We obtained token-level binding site and protein modiﬁcation labels from the Protein Data Bank (Berman et al., 2000).

Original: 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).
Revised: 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 [Figure] were used in model training).

Original: For the diagnostic classiﬁer, we used the respective training splits for training and the validation splits for evaluation.
Revised: For the diagnostic classiﬁer, we used the respective training splits for training and the validation splits for evaluation.

Original: See Appendix B.3 for additional details.
Revised: See Appendix B.4 for additional details.

Original: 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.
Revised: 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.

Revised: We only include results for attention heads where at least 100 high-conﬁdence attention arcs are available for analysis.

Original: We set the attention threshold θ to 0.3 to select for highconﬁdence attention while retaining sufﬁcient data for analysis.
Revised: We set the attention threshold θ to 0. to select for high-conﬁdence attention while retaining sufﬁcient data for analysis.

Original: We truncate all protein sequences to a length of 512 to reduce memory requirements.
Revised: We truncate all protein sequences to a length of 512 to reduce memory requirements.

Original: 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).[Figure]
Revised: 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).

Original:  W HAT DOES ATTENTION UNDERSTAND ABOUT PROTEINS ? 4. P ROTEIN S TRUCTURE
Revised:  W HAT DOES ATTENTION UNDERSTAND ABOUT PROTEINS ? 4. P ROTEIN S TRUCTURE

Original: Here we explore the relationship between attention and tertiary structure, as characterized by contact maps (see Section 2).
Revised: Here we explore the relationship between attention and tertiary structure, as characterized by contact maps (see Section 2).

Original: Secondary structure results are included in Appendix C.1.
Revised: Secondary structure results are included in Appendix C.1.

Original: Attention aligns strongly with contact maps in the deepest layers.
Revised: Attention aligns strongly with contact maps in the deepest layers.

Original: Figure 2 shows the alignment (Eq. 1) between attention and contact maps across  ﬁve models  2 .
Revised: Figure 2 shows how attention aligns with contact maps across the heads of the ﬁve models evaluated 2 , based on the metric deﬁned in Equation 1.

Original: For all models, alignment is strongest in a small number of heads  in the deepest layers (see statistical signiﬁcance tests in Appendix C.2).
Revised: The most aligned heads are found in the deepest layers and focus up to 44.7% (TapeBert), 55.7% (ProtAlbert), 58.5% (ProtBert), 63.2% (ProtBert-BFD), and 44.5% (ProtXLNet) of attention on contacts, whereas the background frequency of contacts among all amino acid pairs in the dataset is 1.3%.

Original: Figure 1a shows an example protein and the induced attention from one such head.
Revised: Figure 1a shows an example protein and the induced attention from the top head in TapeBert.

Revised: Statistical signiﬁcance tests and results for a null model are provided in Appendix C.2.

Revised: [Figure]

Original: Considering the model was trained on a masked language modeling task without any spatial information, the presence of these structurally-aware attention heads is intriguing.
Revised: Considering the model was trained on a masked language modeling task without any spatial information, the presence of these structurally-aware attention heads is intriguing.

Original: One possible reason for this emergent behavior is that contacts are more likely to biochemically interact with one another, creating statistical dependencies between the amino acids in contact.
Revised: One possible reason for this emergent behavior is that contacts are more likely to biochemically interact with one another, creating statistical dependencies between the amino acids in contact.

Original: By focusing attention on the contacts of a masked position, the language model may acquire valuable context for token prediction.
Revised: By focusing attention on the contacts of a masked position, the language model may acquire valuable context for token prediction.

Original: While there seems to be a strong correlation between the attention head output and classically-deﬁned contacts, there are also differences.
Revised: While there seems to be a strong correlation between the attention head output and classically-deﬁned contacts, there are also differences.

Original: The model may have learned a differing contextualized or nuanced formulation that describes amino acid interactions.
Revised: The model may have learned a differing contextualized or nuanced formulation that describes amino acid interactions.

Original: These learned interactions could then be used for further discovery and investigation or repurposed for prediction tasks similar to how principles of coevolution enabled a powerful representation for structure prediction.
Revised: These learned interactions could then be used for further discovery and investigation or repurposed for prediction tasks similar to how principles of coevolution enabled a powerful representation for structure prediction.

Original: Attention is a well-calibrated predictor of contact maps in some cases. It has been implied that attention weights represent a model’s conﬁdence in detecting implicit features (Voita et al., 2019; Correia et al., 2019). To quantify and test this intuition, we propose attention calibration , a measure of the correspondence between attention weight and feature probability.
Revised: 4. B INDING S ITES AND P OST -T RANSLATIONAL M ODIFICATIONS

Original: We base this metric on Expected Calibration Error (Naeini et al., 2015)—or ECE—which estimates the expected difference between a classiﬁer’s conﬁdence and the true class probability through Bayesian binning. We demonstrate this approach by measuring attention calibration with the contact map property f ( i, j ) , a pairwise feature indicating whether the amino acids at positions i and j are in contact.
Revised: We also analyze how attention interacts with binding sites and post-translational modiﬁcations (PTMs), which both play a key role in protein function.

Original: [Figure] [Figure] Especially in the deeper layers, binding sites are targeted at a much higher frequency than would occur by chance (4.8%).
Revised: Attention targets binding sites throughout most layers of the models.

Original: Following the ECE method, we ﬁrst partition all samples (amino-acid pairs ( i, j ) ) from the dataset into 10 equal-sized bins based on conﬁdence (attention weight α i,j ), as shown in the calibration plot in Figure 3 ( x -axis). We then compute the empirical probability of the positive class within each bin, i.e., the proportion of amino acid pairs that are in contact ( y -axis). The ECE is deﬁned as the average absolute difference between empirical probability and conﬁdence (attention) acrossbins. Traditionally a weighted average is used; however, due to the sparsity ofattention, the ﬁrst bucket ( α < 0 .
Revised: Figure 3 shows the proportion of attention focused on binding sites (Eq. 1) across the heads of the 5 models studied.

Original: 1 ) would dominate this calculation. The ﬁrst bin in Figure 3, for example, contains 99.96% ofall examples. Since both attention and empirical probability is low in this bin, a weighted average would produce an artiﬁcially low calibration error. Therefore we instead use an unweighted average.
Revised: Attention to binding sites is most pronounced in the ProtAlbert model (Figure 3b), which has 22 heads that focus over 50% of attention on bindings sites, whereas the background frequency of binding sites in the dataset is 4.8%.

Original: This approach yields a calibration error of 0.09 for the most aligned head from TapeBert (Figure 3) suggesting a moderate level of calibration. The results were mixed for the ProtTrans models  on their respective most-aligned heads: ProtBert (0.13), ProtBert-BFD (0.15), ProtAlbert (0.13), ProtXLNet (0.12). (See Appendix C.2.) However, calibration methods such as isotonic regression (Zadrozny &
Revised: The three BERT models (Figures 3a, 3c, and 3d) also attend strongly to binding sites, with attention heads focusing up to 48.2%, 50.7%, and 45.6% of attention on binding sites, respectively.

Original: Elkan, 2002) may be applied as a post-processing step to improve thecalibration of attention weights and thereby provide a more precise understanding of attention as a measure of conﬁdence. 4. B INDING S ITES AND P OST -T RANSLATIONAL M ODIFICATIONS We also analyze how attentioninteracts with binding sites and post-translational modiﬁcations (PTMs), which both play a key role in protein function. Attention targets binding sites, especially in  the deeper layers.
Revised: Figure 1b visualizes the attention in one strongly-aligned head from the TapeBert model.

Revised: Statistical signiﬁcance tests and a comparison to a null model are provided in Appendix C.3.

Original: As Figure 4 shows, binding sites are attended to most strongly in the ﬁnal six layers of TapeBert, which include 5 heads that focus at least 40% of attention on binding sites. The ProtBert, ProtBert-BFD, and ProtAlbert models showed similar behavior, each having heads that directed up to 51%, 46%, and 57% to binding sites, respectively (see App. C.3).
Revised: ProtXLNet (Figure 3e) also targets binding sites, but not as strongly as the other models: the most aligned head focuses 15.1% of attention on binding sites, and the average head directs just 6.2% of attention to binding sites, compared to 13.2%, 19.8%, 16.0%, and 15.1% for the ﬁrst four models in Figure 3. It’s unclear whether this disparity is due to differences in architectures or pre-training objectives; for example, ProtXLNet uses a bidirectional auto-regressive pretraining method (see Appendix A.1), whereas the other 4 models all use masked language modeling objectives.

Original: We also ﬁnd that binding sites are attended to from far away in the sequence.
Revised: [Figure] [Figure] The embedding probes (orange) quantify the knowledge of the given property that is encoded in each layer’s output embeddings.

Original: In Head 7-1 (Fig. 1b) for example, attention to binding sites spans 124 tokens on average.
Revised: The attention probe (blue), show the amount of information encoded in attention weights for the (pairwise) contact feature.

Revised: Additional details are provided in Appendix B.3.

Original: Why does attention target binding sites?
Revised: Why does attention target binding sites?

Original: In contrast to contact maps, which reveal relationships within proteins, binding sites describe how a protein interacts with other molecules.
Revised: In contrast to contact maps, which reveal relationships within proteins, binding sites describe how a protein interacts with other molecules.

Original: These external interactions ultimately deﬁne the high-level function of the protein, and thus binding sites remain conserved even when the sequence as a whole evolves (Kinjo & Nakamura, 2009).
Revised: These external interactions ultimately deﬁne the high-level function of the protein, and thus binding sites remain conserved even when the sequence as a whole evolves (Kinjo & Nakamura, 2009).

Original: Further, structural motifs in binding sites are mainly restricted to speciﬁc families or superfamilies of proteins (Kinjo &
Revised: Further, structural motifs in binding sites are mainly restricted to speciﬁc families or superfamilies of proteins (Kinjo &

Original: Nakamura, 2009), and binding sites can reveal evolutionary relationships among proteins (Lee et al., 2017).
Revised: Nakamura, 2009), and binding sites can reveal evolutionary relationships among proteins (Lee et al., 2017).

Original: Thus binding sites may provide the model with a high-level characterization of the protein that is robust to individual sequence variation.
Revised: Thus binding sites may provide the model with a high-level characterization of the protein that is robust to individual sequence variation.

Revised: By attending to these regions, the model can leverage this higher-level context when predicting masked tokens throughout the sequence.

Original: Attention targets PTMs in a small number of heads.
Revised: Attention targets PTMs in a small number of heads.

Original: A small number of heads in each model concentrate their attention very strongly on  post-translational modiﬁcations (PTMs).
Revised: A small number of heads in each model concentrate their attention very strongly on amino acids associated with post-translational modiﬁcations (PTMs).

Original: For example, Head 11-6 in TapeBert focused 64% of attention on PTMs, though PTMs occur at only 0.8% of [Figure]
Revised: For example, Head 11-6 in TapeBert focused 64% of attention on PTM positions, though these occur at only 0.8% of sequence positions in the dataset.

Original: [Figure] sequence positions in the dataset. 4 Similar to our discussion on binding sites, PTMs are critical to protein function (Rubin & Rosen, 1975) and thereby are likely to exhibit behavior that is conserved across the sequence space.
Revised: 3 Similar to our discussion on binding sites, PTMs are critical to protein function (Rubin & Rosen, 1975) and thereby are likely to exhibit behavior that is conserved across the sequence space.

Original: See Appendix C.4 for full results.
Revised: See Appendix C.4 for full results.

Original: 4. C ROSS -L AYER A NALYSIS
Revised: 4. C ROSS -L AYER A NALYSIS

Original: We analyze how attention captures properties of varying complexity across different layers of the model, and compare to a probing analysis of embeddings and attention weights (see Section 3).
Revised: We analyze how attention captures properties of varying complexity across different layers of TapeBert, and compare to a probing analysis of embeddings and attention weights (see Section 3).

Original: Attention targets higher-level properties in deeper layers.
Revised: Attention targets higher-level properties in deeper layers.

Original: As shown in Figure 5, deeper layers focus relatively more attention on binding sites and contacts (high-level concept), whereas secondary structure (low- to mid-level concept) is targeted more evenly across layers. 
Revised: As shown in Figure 4, deeper layers focus relatively more attention on binding sites and contacts (high-level concept), whereas secondary structure (low- to mid-level concept) is targeted more evenly across layers. The probing analysis [Figure] [Figure]

Revised: QRSTVWY

Original: The probing analysis of attention (Figure 6, blue) similarly shows that knowledge of contact maps (a pairwise feature) is encoded in attention weights primarily in the last 1-2 layers.
Revised: [Figure] [Figure] of attention (Figure 5, blue) similarly shows that knowledge of contact maps (a pairwise feature) is encoded in attention weights primarily in the last 1-2 layers.

Original: These results are consistent with prior work in NLP that suggests deeper layers in text-based Transformers attend to more complex properties (Vig & Belinkov, 2019) and encode higher-level representations (Raganato & Tiedemann, 2018; Peters et al., 2018; Tenney et al., 2019; Jawahar et al., 2019).
Revised: These results are consistent with prior work in NLP that suggests deeper layers in text-based Transformers attend to more complex properties (Vig & Belinkov, 2019) and encode higher-level representations (Raganato & Tiedemann, 2018; Peters et al., 2018; Tenney et al., 2019; Jawahar et al., 2019).

Original: The embedding probes (Figure 6, orange) also show that the model ﬁrst builds representations of local secondary structure in lower layers before fully encoding binding sites and contact maps in deeper layers.
Revised: The embedding probes (Figure 5, orange) also show that the model ﬁrst builds representations of local secondary structure in lower layers before fully encoding binding sites and contact maps in deeper layers.

Original: However, this analysis also reveals stark differences in how knowledge of contact maps is accrued in embeddings, which accumulate this knowledge gradually over many layers, compared to attention weights, which acquire this knowledge only in the ﬁnal layers in this case.
Revised: However, this analysis also reveals stark differences in how knowledge of contact maps is accrued in embeddings, which accumulate this knowledge gradually over many layers, compared to attention weights, which acquire this knowledge only in the ﬁnal layers in this case.

Original: This example points out limitations of common layerwise probing approaches that only consider embeddings, which, intuitively, represent what the model knows but not necessarily how it operationalizes that knowledge.[-2 4 6 8 10 12 Head   
 L a y e r % Attention 100% Max 0% 20% 40% 60% 80% 2 4 6 8 10 12 Head   
-]
Revised: This example points out limitations of common layerwise probing approaches that only consider embeddings, which, intuitively, represent what the model knows but not necessarily how it operationalizes that knowledge.[-2 4 6 8 10 12 Head   
 L a y e r % Attention 100% Max 0% 20% 40% 60% 80% 2 4 6 8 10 12 Head   
-]

Original: L a y e r % Attention 20%

Original: Max 0% 5% 10% 15% 20%

Original: A C D E F G H I K L M N P Q R S T V W Y

Original: ACDEFGHIKLMNPQRST

Original: V WY 0.

Original: A C D E F G H I K L M N P Q R S T V W Y

Original: ACDEFGHIKLMNPQRST

Original: V WY

Original: 4. A MINO A CIDS AND THE S UBSTITUTION M ATRIX
Revised: 4. A MINO A CIDS AND THE S UBSTITUTION M ATRIX

Original: In addition to high-level structural and functional properties, we also performed a ﬁne-grained analysis of the interaction between attention and particular amino acids.
Revised: In addition to high-level structural and functional properties, we also performed a ﬁne-grained analysis of the interaction between attention and particular amino acids.

Original: Attention heads specialize in particular amino acids.
Revised: Attention heads specialize in particular amino acids.

Original: We computed the proportion of TapeBert’s attention to each of the 20 standard amino acids, as shown in Figure 7 for two example amino acids.
Revised: We computed the proportion of TapeBert’s attention to each of the 20 standard amino acids, as shown in Figure 6 for two example amino acids.

Original: For 16 of the amino acids, there exists an attention head that focuses over 25% of attention on that amino acid, signiﬁcantly greater than the background frequencies of the corresponding amino acids, which range from 1.3% to 9.4% (see App.
Revised: For 16 of the amino acids, there exists an attention head that focuses over 25% of attention on that amino acid, signiﬁcantly greater than the background frequencies of the corresponding amino acids, which range from 1.3% to 9.4%.

Original: C.5 for detailed results). Similar behavior was observed for ProtBert, ProtBert-BFD, ProtAlbert, and ProtXLNet models, with 17, 15, 16, and 18 amino acids, respectively, receiving greater than 25% of the attention from at least one attention head. 
Revised: Similar behavior was observed for ProtBert, ProtBert-BFD, ProtAlbert, and ProtXLNet models, with 17, 15, 16, and 18 amino acids, respectively, receiving greater than 25% of the attention from at least one attention head. Detailed results for TapeBert including statistical signiﬁcance tests and comparison to a null model are presented in Appendix C.5.

Original: Attention is consistent with substitution relationships.
Revised: Attention is consistent with substitution relationships.

Original: A natural follow-up question  is whether each head has “memorized” speciﬁc amino acids to target, or whether it has actually learned meaningful properties that correlate with particular amino acids.
Revised: A natural follow-up question from the above analysis is whether each head has “memorized” speciﬁc amino acids to target, or whether it has actually learned meaningful properties that correlate with particular amino acids.

Original: To test the latter hypothesis, we analyze whether amino acids with similar structural and functional properties are attended to similarly across heads.
Revised: To test the latter hypothesis, we analyze whether amino acids with similar structural and functional properties are attended to similarly across heads.

Original: Speciﬁcally, we compute the Pearson correlation between the distribution of attention across heads between all pairs of distinct amino acids, as shown in Figure 8 (left).
Revised: Speciﬁcally, we compute the Pearson correlation between the distribution of attention across heads between all pairs of distinct amino acids, as shown in Figure (left) for TapeBert.

Original: For example, the entry for Pro (P) and Phe (F) is the correlation between the two heatmaps in Figure 7.
Revised: For example, the entry for Pro (P) and Phe (F) is the correlation between the two heatmaps in Figure 6.

Original: We compare these scores to the BLOSUM62 substitution scores (Sec. 2) in Figure 8 (right), and ﬁnd a Pearson correlation of 0.73, suggesting that attention is moderately consistent with substitution relationships.
Revised: We compare these scores to the BLOSUM62 substitution scores (Sec. 2) in Figure 7 (right), and ﬁnd a Pearson correlation of 0.73, suggesting that attention is moderately consistent with substitution relationships.

Original: Similar correlations are observed for ProtBert (0.68), ProtBert-BFD (0.75), ProtXLNet (0.71), and 
Revised: Similar correlations are observed for the ProtTrans models: 0.68 (ProtBert), 0.75 (ProtBert-BFD), 0.60 (ProtAlbert), and 0.71 (ProtXLNet).

Original: ProtAlbert (0.60).
Revised: As a baseline, the randomized versions of these models (Appendix B.2) yielded correlations of -0.02 (TapeBert), 0. (ProtBert), -0.03 (ProtBert-BFD), -0.05 (ProtAlbert), and 0.21 (ProtXLNet).

Original:  R ELATED W ORK 5. P ROTEIN LANGUAGE MODELS
Revised:  R ELATED W ORK 5. P ROTEIN LANGUAGE MODELS

Original: Deep neural networks for protein language modeling have received broad interest.
Revised: Deep neural networks for protein language modeling have received broad interest.

Original: Early work applied the Skip-gram model (Mikolov et al., 2013) to construct continuous embeddings from protein sequences (Asgari & Mofrad, 2015).
Revised: Early work applied the Skip-gram model (Mikolov et al., 2013) to construct continuous embeddings from protein sequences (Asgari & Mofrad, 2015).

Original: Sequence-only language models have since been trained through autoregressive or autoencoding self-supervision objectives for discriminative and generative tasks, for example, using LSTMs or Transformer-based architectures (Alley et al., 2019; Bepler & Berger, 2019;
Revised: Sequence-only language models have since been trained through autoregressive or autoencoding self-supervision objectives for discriminative and generative tasks, for example, using LSTMs or Transformer-based architectures (Alley et al., 2019; Bepler & Berger, 2019;

Original: Rao et al., 2019; Rives et al., 2019).
Revised: Rao et al., 2019; Rives et al., 2019).

Original: TAPE created a benchmark of ﬁve tasks to assess protein sequence models, and ProtTrans also released several large-scale pretrained protein Transformer models (Elnaggar et al., 2020).
Revised: TAPE created a benchmark of ﬁve tasks to assess protein sequence models, and ProtTrans also released several large-scale pretrained protein Transformer models (Elnaggar et al., 2020).

Original: Riesselman et al.
Revised: Riesselman et al.

Original: ; Madani et al.
Revised: ; Madani et al.

Original: (2020) trained autoregressive generative models to predict the functional effect of mutations and generate natural-like proteins.
Revised: (2020) trained autoregressive generative models to predict the functional effect of mutations and generate natural-like proteins.

Original: From an interpretability perspective, Rives et al.
Revised: From an interpretability perspective, Rives et al.

Original: (2019) showed that the output embeddings from a pretrained Transformer can recapitulate structural and functional properties of proteins through learned linear transformations.
Revised: (2019) showed that the output embeddings from a pretrained Transformer can recapitulate structural and functional properties of proteins through learned linear transformations.

Revised: Various works have analyzed output embeddings of protein models through dimensionality reduction techniques such as PCA or t-SNE (Elnaggar et al., 2020; Biswas et al., 2020).

Original: In our work, we take an interpretability-ﬁrst perspective to focus on the internal model representations, speciﬁcally attention and intermediate hidden states, across multiple protein language models.
Revised: In our work, we take an interpretability-ﬁrst perspective to focus on the internal model representations, speciﬁcally attention and intermediate hidden states, across multiple protein language models.

Original: We also explore novel biological properties including binding sites and post-translational modiﬁcations.
Revised: We also explore novel biological properties including binding sites and post-translational modiﬁcations.

Original: 5. I NTERPRETING M ODELS IN NLP
Revised: 5. I NTERPRETING M ODELS IN NLP

Original: The rise of deep neural networks in ML has also led to much work on interpreting these so-called black-box models.
Revised: The rise of deep neural networks in ML has also led to much work on interpreting these so-called black-box models.

Original: This section reviews the NLP interpretability literature on the Transformer model, which is directly comparable to our work on interpreting Transformer models of protein sequences.
Revised: This section reviews the NLP interpretability literature on the Transformer model, which is directly comparable to our work on interpreting Transformer models of protein sequences.

Original: Interpreting Transformers.
Revised: Interpreting Transformers.

Original: The Transformer is a neural architecture that uses attention to accelerate learning (Vaswani et al., 2017).
Revised: The Transformer is a neural architecture that uses attention to accelerate learning (Vaswani et al., 2017).

Original: In NLP, transformers are the backbone of state-of-the-art pre-trained language models such as BERT (Devlin et al., 2019).
Revised: In NLP, transformers are the backbone of state-of-the-art pre-trained language models such as BERT (Devlin et al., 2019).

Original: BERTology focuses on interpreting what the BERT model learns about language using a suite of probes and interventions (Rogers et al., 2020).
Revised: BERTology focuses on interpreting what the BERT model learns about language using a suite of probes and interventions (Rogers et al., 2020).

Original: So-called diagnostic classiﬁers are used to interpret the outputs from BERT’s layers (Veldhoen et al., 2016).
Revised: So-called diagnostic classiﬁers are used to interpret the outputs from BERT’s layers (Veldhoen et al., 2016).

Original: At a high level, mechanisms for interpreting BERT can be placed into three main categories: interpreting the learned embeddings (Ethayarajh, 2019; Wiedemann et al., 2019; Mickus et al., 2019; Adi et al., 2016; Conneau et al., 2018), BERT’s learned knowledge of syntax (Lin et al., 2019; Liu et al., 2019; Tenney et al., 2019; Htut et al., 2019; Hewitt & Manning, 2019; Goldberg, 2019), and BERT’s learned knowledge of semantics (Tenney et al., 2019;
Revised: At a high level, mechanisms for interpreting BERT can be placed into three main categories: interpreting the learned embeddings (Ethayarajh, 2019; Wiedemann et al., 2019; Mickus et al., 2019; Adi et al., 2016; Conneau et al., 2018), BERT’s learned knowledge of syntax (Lin et al., 2019; Liu et al., 2019; Tenney et al., 2019; Htut et al., 2019; Hewitt & Manning, 2019; Goldberg, 2019), and BERT’s learned knowledge of semantics (Tenney et al., 2019;

Original: Ettinger, 2020).
Revised: Ettinger, 2020).

Original: Interpreting attention speciﬁcally.
Revised: Interpreting attention speciﬁcally.

Original: Interpreting attention on textual sequences is a wellestablished area of research (Wiegreffe & Pinter, 2019; Zhong et al., 2019;
Revised: Interpreting attention on textual sequences is a wellestablished area of research (Wiegreffe & Pinter, 2019; Zhong et al., 2019;

Original: Brunner et al., 2020;
Revised: Brunner et al., 2020;

Original: Hewitt & Manning, 2019).
Revised: Hewitt & Manning, 2019).

Original: Past work has been shown that attention correlates with syntactic and semantic relationships in natural language in some cases (Clark et al., 2019; Vig & Belinkov, 2019;
Revised: Past work has been shown that attention correlates with syntactic and semantic relationships in natural language in some cases (Clark et al., 2019; Vig & Belinkov, 2019;

Original: Htut et al., 2019).
Revised: Htut et al., 2019).

Original: Depending on the task and model architecture, attention may have less or more explanatory power for model predictions (Jain & Wallace, 2019; Serrano & Smith, 2019; Pruthi et al., 2020; Moradi et al., 2019; Vashishth et al., 2019).
Revised: Depending on the task and model architecture, attention may have less or more explanatory power for model predictions (Jain & Wallace, 2019; Serrano & Smith, 2019; Pruthi et al., 2020; Moradi et al., 2019; Vashishth et al., 2019).

Original: Visualization techniques have been used to convey the structure and properties of attention in Transformers (Vaswani et al., 2017; Kovaleva et al., 2019;
Revised: Visualization techniques have been used to convey the structure and properties of attention in Transformers (Vaswani et al., 2017; Kovaleva et al., 2019;

Original: Hoover et al., 2019; Vig, 2019).
Revised: Hoover et al., 2019; Vig, 2019).

Original: Recent work has begun to analyze attention in Transformer models outside of the domain of natural language (Schwaller et al., 2020; Payne et al., 2020).
Revised: Recent work has begun to analyze attention in Transformer models outside of the domain of natural language (Schwaller et al., 2020; Payne et al., 2020).

Original: Our work extends these methods to protein sequence models, but also introduces methods for interpreting Transformers that, to our knowledge, have not previously been applied in NLP or other domains. First, we introduce the notion of attention calibration , the degree to which attention weights correspond to the actual probability of particular features.
Revised: Our work extends these methods to protein sequence models by considering particular biophysical properties and relationships.

Original: Second, we present  cross-layer probing analyses ofboth attention weights and layer embeddings.
Revised: We also present a joint cross-layer probing analysis of attention weights and layer embeddings.

Original: While past work  has analyzed attention and embeddings across layers, we believe we are the ﬁrst to do so  using a single, uniﬁed metric, which enables us to directly compare the relative information content of the two representations. 
Revised: While past work in NLP has analyzed attention and embeddings across layers, we believe we are the ﬁrst to do so in any domain using a single, uniﬁed metric, which enables us to directly compare the relative information content of the two representations. Finally, we present a novel tool for visualizing attention embedded in three-dimensional structure.

Original:  C ONCLUSIONS AND F UTURE W ORK
Revised:  C ONCLUSIONS AND F UTURE W ORK

Original: This paper builds on the synergy between NLP and computational biology by adapting and extending
Revised: This paper builds on the synergy between NLP and computational biology by adapting and extending

Original: NLP interpretability methods to protein sequence modeling.
Revised: NLP interpretability methods to protein sequence modeling.

Original: We show how a Transformer language model recovers structural and functional properties of proteins and integrates this knowledge directly into its attention mechanism.
Revised: We show how a Transformer language model recovers structural and functional properties of proteins and integrates this knowledge directly into its attention mechanism.

Original: In some cases, attention also provides a well-calibrated measure of conﬁdence in the model’s predictions of these properties.

Original: While this paper focuses on reconciling attention with known properties of proteins, one might also leverage attention to uncover novel relationships or more nuanced forms of existing measures such as contacts maps, as discussed in Section 4.1.
Revised: While this paper focuses on reconciling attention with known properties of proteins, one might also leverage attention to uncover novel relationships or more nuanced forms of existing measures such as contacts maps, as discussed in Section 4.1.

Original: In this way, language models have the potential to serve as tools for scientiﬁc discovery.
Revised: In this way, language models have the potential to serve as tools for scientiﬁc discovery.

Original: But in order for learned representations to be accessible to domain experts, they must be presented in an appropriate context to facilitate discovery.
Revised: But in order for learned representations to be accessible to domain experts, they must be presented in an appropriate context to facilitate discovery.

Original: Visualizing attention in the context of protein structure (Figure 1) is one attempt to do so.
Revised: Visualizing attention in the context of protein structure (Figure 1) is one attempt to do so.

Original: We believe there is the potential to develop such contextual visualizations of learned representations in a range of scientiﬁc domains.
Revised: We believe there is the potential to develop such contextual visualizations of learned representations in a range of scientiﬁc domains.

Original: A M ODEL O VERVIEW
Revised: A M ODEL O VERVIEW

Original: A.
Revised: A.

Original: BERT T RANSFORMER
Revised: BERT T RANSFORMER

Original: A RCHITECTURE
Revised: A RCHITECTURE

Original: Stacked Encoder: BERT uses a stacked-encoder architecture, which inputs a sequence of tokens x
Revised: Stacked Encoder: BERT uses a stacked-encoder architecture, which inputs a sequence of tokens x

Original: = ( x 1 , ..., x n ) and applies position and token embeddings followed by a series of encoder layers.
Revised: = ( x 1 , ..., x n ) and applies position and token embeddings followed by a series of encoder layers.

Original: Each layer applies multi-head self-attention (see below) in combination with a feedforward network, layer normalization, and residual connections.
Revised: Each layer applies multi-head self-attention (see below) in combination with a feedforward network, layer normalization, and residual connections.

Original: The output of each layer (cid:96) is a sequence of contextualized embeddings ( h ( (cid:96) ) 1 , . .
Revised: The output of each layer (cid:96) is a sequence of contextualized embeddings ( h ( (cid:96) ) 1 , . .

Original: , h ( (cid:96) ) n ) .
Revised: , h ( (cid:96) ) n ) .

Original: Self-Attention:
Revised: Self-Attention:

Original: Given an input x = ( x 1 , . . . , x n ) , the self-attention mechanism assigns to each token pair i , j an attention weight α i,j > 0 where (cid:80) j α i,j = 1 .
Revised: Given an input x = ( x 1 , . . . , x n ) , the self-attention mechanism assigns to each token pair i , j an attention weight α i,j > 0 where (cid:80) j α i,j = 1 .

Original: Attention in BERT is bidirectional.
Revised: Attention in BERT is bidirectional.

Original: In the multi-layer, multi-head setting, α is speciﬁc to a layer and head.
Revised: In the multi-layer, multi-head setting, α is speciﬁc to a layer and head.

Original: The BERT-Base model haslayers and 12 heads.
Revised: The BERT-Base model haslayers and 12 heads.

Original: Each attention head learns a distinct set of weights, resulting in 12 x 12 =distinct attention mechanisms in this case.
Revised: Each attention head learns a distinct set of weights, resulting in 12 x 12 =distinct attention mechanisms in this case.

Original: The attention weights α i,j are computed from the scaled dot-product of the query vector of i and the key vector of j , followed by a softmax operation.
Revised: The attention weights α i,j are computed from the scaled dot-product of the query vector of i and the key vector of j , followed by a softmax operation.

Original: The attention weights are then used to produce a weighted sum of value vectors:
Revised: The attention weights are then used to produce a weighted sum of value vectors:

Original: [Equation]
Revised: [Equation]

Original: using query matrix Q , key matrix K , and value matrix V , where d k is the dimension of K .
Revised: using query matrix Q , key matrix K , and value matrix V , where d k is the dimension of K .

Original: In a multi-head setting, the queries, keys, and values are linearly projected h times, and the attention operation is performed in parallel for each representation, with the results concatenated.
Revised: In a multi-head setting, the queries, keys, and values are linearly projected h times, and the attention operation is performed in parallel for each representation, with the results concatenated.

Original: A. O THER T RANSFORMER V ARIANTS
Revised: A. O THER T RANSFORMER V ARIANTS

Original: ALBERT:
Revised: ALBERT:

Original: The architecture of ALBERT differs from BERT in two ways: (1) It shares parameters across layers, unlike BERT which learns distinct parameters for every layer and (2) It uses factorized embeddings, which allows the input token embeddings to be of a different (smaller) size than the hidden states.
Revised: The architecture of ALBERT differs from BERT in two ways: (1) It shares parameters across layers, unlike BERT which learns distinct parameters for every layer and (2) It uses factorized embeddings, which allows the input token embeddings to be of a different (smaller) size than the hidden states.

Original: The original version of ALBERT designed for text also employed a sentence-order prediction pretraining task, but this was not used on the models studied in this paper.
Revised: The original version of ALBERT designed for text also employed a sentence-order prediction pretraining task, but this was not used on the models studied in this paper.

Original: XLNet: Instead of the masked-language modeling pretraining objective use for BERT, XLNet uses a bidirectional auto-regressive pretraining method that considers all possible orderings of the input factorization.
Revised: XLNet: Instead of the masked-language modeling pretraining objective use for BERT, XLNet uses a bidirectional auto-regressive pretraining method that considers all possible orderings of the input factorization.

Original: The architecture also adds a segment recurrence mechanism to process long sequences, as well as a relative rather than absolute encoding scheme.
Revised: The architecture also adds a segment recurrence mechanism to process long sequences, as well as a relative rather than absolute encoding scheme.

Original: B A DDITIONAL E XPERIMENTAL D ETAILS B. A LTERNATIVE ATTENTION ALIGNMENT METRIC
Revised: B A DDITIONAL E XPERIMENTAL D ETAILS B. A LTERNATIVE ATTENTION AGREEMENT METRIC

Original: Here we present an alternative formulation to Eq.
Revised: Here we present an alternative formulation to Eq.

Original: 1 based on an attention-weighted average.
Revised: 1 based on an attention-weighted average.

Original: We deﬁne an indicator function f
Revised: We deﬁne an indicator function f

Original: ( i, j ) for property f that returns 1 if the property is present in token pair ( i, j ) (i.e., if amino acids i and j are in contact), and zero otherwise.
Revised: ( i, j ) for property f that returns 1 if the property is present in token pair ( i, j ) (i.e., if amino acids i and j are in contact), and zero otherwise.

Original: We then compute the proportion of attention that aligns with f over a dataset X as follows:
Revised: We then compute the proportion of attention that matches with f over a dataset X as follows:

Original: [Equation]
Revised: [Equation]

Original: where α i,j ( x ) denotes the attention from i to j for input sequence x .
Revised: where α i,j ( x ) denotes the attention from i to j for input sequence x .

Original: B. P ROBING METHODOLOGY
Revised: B. S TATISTICAL SIGNIFICANCE TESTING AND NULL MODELS

Revised: We perform statistical signiﬁcance tests to determine whether any results based on the metric deﬁned in Equation 1 are due to chance. Given a property f , as deﬁned in Section 3, we perform a twoproportion z-test comparing (1) the proportion of high-conﬁdence attention arcs ( α i,j > θ ) for which f ( i, j )

Revised: = 1 , and (2) the proportion of all possible pairs i, j for which f

Revised: ( i, j )

Revised: = 1 .

Revised: Note that the ﬁrst proportion is exactly the metric p α ( f ) deﬁned in Equation 1 (e.g. the proportion of attention aligned with contact maps).

Revised: The second proportion is simply the background frequency of the property (e.g. the background frequency of contacts).

Revised: Since we extract the maximum scores over all of the heads in the model, we treat this as a case of multiple hypothesis testing and apply the Bonferroni correction, with the number of hypotheses m equal to the number of attention heads.

Revised: As an additional check that the results did not occur by chance, we also report results on baseline (null) models.

Revised: We initially considered using randomly initialized versions of the models studied as null models.

Revised: However, none of these models yielded attention weights greater than the attention threshold θ .

Revised: This suggests that the mere existence of the high-conﬁdence attention weights used in the analysis could not have occurred by chance but it does not shed light on the particular analyses performed.

Revised: Therefore, we implemented an alternative randomization scheme in which we randomly shufﬂe attention weights from the original models as a post-processing step.

Revised: Speciﬁcally, we permute the sequence of attention weights from each token for a given attention head.

Revised: To illustrate, let’s say that the original model produced attention weights of (0.3, 0.2, 0.1, 0.4, 0.0) from position i in protein sequence x from head h , where | x | = 5 .

Revised: In the null model, the attention weights from position i in sequence x in head h would be a random permutation of those weights, e.g., (0.2, 0.0, 0.4, 0.3, 0.1).

Revised: Note that these are still valid attention weights as they would sum to 1 (since the original weights would sum to 1 by deﬁnition).

Revised: We report results using this form of baseline model.

Revised: B. P ROBING METHODOLOGY

Original: Embedding probe.
Revised: Embedding probe.

Original: We probe the embedding vectors output from each layer using a linear probing classiﬁer.
Revised: We probe the embedding vectors output from each layer using a linear probing classiﬁer.

Original: For token-level probing tasks (binding sites, secondary structure) we feed each token’s output vector directly to the classiﬁer.
Revised: For token-level probing tasks (binding sites, secondary structure) we feed each token’s output vector directly to the classiﬁer.

Original: For token-pair probing tasks (contact map) we construct a pairwise feature vector by concatenating the elementwise differences and products of the two tokens’ output vectors, following the TAPE 5 implementation.
Revised: For token-pair probing tasks (contact map) we construct a pairwise feature vector by concatenating the elementwise differences and products of the two tokens’ output vectors, following the TAPE 4 implementation.

Original: We use task-speciﬁc evaluation metrics for the probing classiﬁer: for secondary structure prediction, we measure F1 score; for contact prediction, we measure precision@ L/ 5 , where L is the length of the protein sequence, following standard practice (Moult et al., 2018); for binding site prediction, we measure precision@ L/ 20 , since approximately one in twenty amino acids in each sequence is a binding site (4.8% in the dataset).
Revised: We use task-speciﬁc evaluation metrics for the probing classiﬁer: for secondary structure prediction, we measure F1 score; for contact prediction, we measure precision@ L/ 5 , where L is the length of the protein sequence, following standard practice (Moult et al., 2018); for binding site prediction, we measure precision@ L/ 20 , since approximately one in twenty amino acids in each sequence is a binding site (4.8% in the dataset).

Original: Attention probe.
Revised: Attention probe.

Original: Just as the attention weight α i,j is deﬁned for a pair of amino acids ( i, j ) , so is the contact property f
Revised: Just as the attention weight α i,j is deﬁned for a pair of amino acids ( i, j ) , so is the contact property f

Original: ( i, j ) , which returns true if amino acids i and j are in contact.
Revised: ( i, j ) , which returns true if amino acids i and j are in contact.

Original: Treating the attention weight as a feature of a token-pair ( i, j ) , we can train a probing classiﬁer that predicts the contact property based on this feature, thereby quantifying the attention mechanism’s knowledge of that property.
Revised: Treating the attention weight as a feature of a token-pair ( i, j ) , we can train a probing classiﬁer that predicts the contact property based on this feature, thereby quantifying the attention mechanism’s knowledge of that property.

Original: In our multi-head setting, we treat the attention weights across all heads in a given layer as a feature vector, and use a probing classiﬁer to assess the knowledge of a given property in the attention weights across the entire layer.
Revised: In our multi-head setting, we treat the attention weights across all heads in a given layer as a feature vector, and use a probing classiﬁer to assess the knowledge of a given property in the attention weights across the entire layer.

Original: As with the embedding probe, we measure performance of the probing classiﬁer using precision@ L/ 5 , where L is the length of the protein sequence, following standard practice for contact prediction.
Revised: As with the embedding probe, we measure performance of the probing classiﬁer using precision@ L/ 5 , where L is the length of the protein sequence, following standard practice for contact prediction.

Original: B. D ATASETS
Revised: B. D ATASETS

Original: We used two protein sequence datasets from the TAPE repository for the analysis: the ProteinNet dataset (AlQuraishi, 2019;
Revised: We used two protein sequence datasets from the TAPE repository for the analysis: the ProteinNet dataset (AlQuraishi, 2019;

Original: Fox et al., 2013; Berman et al., 2000; Moult et al., 2018) and the Secondary
Revised: Fox et al., 2013; Berman et al., 2000; Moult et al., 2018) and the Secondary

Original: Structure dataset (Rao et al., 2019; Berman et al., 2000; Moult et al., 2018; Klausen et al., 2019).
Revised: Structure dataset (Rao et al., 2019; Berman et al., 2000; Moult et al., 2018; Klausen et al., 2019).

Original: The former was used for analysis of amino acids and contact maps, and the latter was used for analysis of secondary structure.
Revised: The former was used for analysis of amino acids and contact maps, and the latter was used for analysis of secondary structure.

Original: We additionally created a third dataset for binding site and post-translational modiﬁcation (PTM) analysis from the Secondary Structure dataset, which was augmented with binding site and PTM annotations obtained from the Protein Data Bank’s Web API.
Revised: We additionally created a third dataset for binding site and post-translational modiﬁcation (PTM) analysis from the Secondary Structure dataset, which was augmented with binding site and PTM annotations obtained from the Protein Data Bank’s Web API.

Original: We excluded any sequences for which annotations were not available.
Revised: We excluded any sequences for which annotations were not available.

Original: The resulting dataset sizes are shown in
Revised: The resulting dataset sizes are shown in

Original: Table 1.
Revised: Table 1.

Original: For the analysis of attention, a random subset of 5000 sequences from the training split of each dataset was used, as the analysis was purely evaluative.
Revised: For the analysis of attention, a random subset of 5000 sequences from the training split of each dataset was used, as the analysis was purely evaluative.

Original: For training and evaluating the diagnostic classiﬁer, the full training and validation splits were used.
Revised: For training and evaluating the diagnostic classiﬁer, the full training and validation splits were used [Figure] [Table] [Figure]

Original: [Table]
Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Table]

Original: C A DDITIONAL R ESULTS OF A TTENTION A NALYSIS C. S ECONDARY S TRUCTURE
Revised: C A DDITIONAL R ESULTS OF A TTENTION A NALYSIS C. S ECONDARY S TRUCTURE

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Revised: C. C ONTACT M APS : S TATISTICAL SIGNIFICANCE TESTS AND NULL MODELS

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Revised: A tt e n t i o n %

Original: [Figure]
Revised: [Figure]

Revised: i o

Original: [Figure]
Revised: [Figure]

Revised: C. B INDING S ITES : S TATISTICAL SIGNIFICANCE TESTS AND NULL MODEL

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Revised: t i o

Original: [Figure]
Revised: [Figure]

Original: INDING S ITES
Revised: C. P OST -T RANSLATIONAL M ODIFICATIONS (PTM S )

Original: [Figure]
Revised: [Figure]

Original: [Figure] 
Revised: [Figure] [Figure] [Figure]

Original: [Figure] 
Revised: [Figure] [Table] Arg R

Original: [Figure]
Revised: Arginine 5. Asn N Asparagine 4.

Original: [Figure]
Revised: Asp D Aspartic acid 5. Cys

Original: [Figure]
Revised: C Cysteine 1. Gln Q

Original: [Table]
Revised: Glutamine 3. Glu E Glutamic acid 6.

Original: Gly
Revised: Gly

Original: G
Revised: G

Original: Glycine 298.
Revised: Glycine 7.

Original: His
Revised: His

Original: H
Revised: H

Original: Histidine 956.
Revised: Histidine 2.

Original: Ile
Revised: Ile

Original: I
Revised: I

Original: Isoleucine 1127.
Revised: Isoleucine 5.

Original: Leu
Revised: Leu

Original: L
Revised: L

Original: Leucine 244.
Revised: Leucine 9.

Original: Lys
Revised: Lys

Original: K
Revised: K

Original: Lysine 1229.
Revised: Lysine 6.

Original: Met
Revised: Met

Original: M
Revised: M

Original: Methionine 373.
Revised: Methionine 2.

Original: Phe
Revised: Phe

Original: F
Revised: F

Original: Phenylalanine 1222.
Revised: Phenylalanine 3.

Original: Pro
Revised: Pro

Original: P
Revised: P

Original: Proline 198.
Revised: Proline 4.

Original: Ser
Revised: Ser

Original: S
Revised: S

Original: Serine 1236.
Revised: Serine 6.

Original: Thr
Revised: Thr

Original: T
Revised: T

Original: Threonine 1219.
Revised: Threonine 5.

Original: Trp
Revised: Trp

Original: W
Revised: W

Original: Tryptophan 1168.
Revised: Tryptophan 1.

Original: Tyr
Revised: Tyr

Original: Y
Revised: Y

Original: Tyrosine 1251.
Revised: Tyrosine 3. [Table]

Revised: D A TTENTION C ALIBRATION A NALYSIS

Original: Val
Revised: [Figure] [Table] [Table] [Table]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: It has been implied that attention weights represent a model’s conﬁdence in detecting particular properties (Voita et al., 2019; Correia et al., 2019).

Revised: To quantify and test this intuition, we propose attention calibration , a measure of the correspondence between attention weight and feature probability.

Revised: We base this metric on Expected Calibration Error (Naeini et al., 2015)—or ECE—which estimates the expected difference between a classiﬁer’s conﬁdence and the true class probability through Bayesian binning.

Original: V
Revised: We demonstrate this approach by measuring attention calibration with the contact map property f ( i, j ) , a pairwise feature indicating whether the amino acids at positions i and j are in contact. Following the ECE method, we ﬁrst partition all samples (amino-acid pairs ( i, j ) ) from the dataset into 10 equal-sized bins based on conﬁdence (attention weight α i,j ), as shown in the calibration plots in Figure 20 ( x -axis). We then compute the empirical probability of the positive class within each bin, i.e., the proportion of amino acid pairs that are in contact ( y -axis).

Revised: The ECE is deﬁned as the average absolute difference between empirical probability and conﬁdence (attention) across bins.

Revised: Traditionally a weighted average is used; however, due to the sparsity of attention, the ﬁrst bucket ( α < 0 .

Original: Valine 1234.
Revised: 1 ) would dominate this calculation. The ﬁrst bin in Figure 20a, for example, contains 99.96% of all examples. Since both attention and empirical probability is low in this bin, a weighted average would produce an artiﬁcially low calibration error. Therefore we instead use an unweighted average.

Revised: This approach yields a calibration error of 0.09 for the most contact-focused head from TapeBert (Figure 20a) suggesting a moderate level of calibration.

Revised: The errors were higher for the the ProtTrans models on their respective most contact-focused heads: ProtBert (0.13), ProtBert-BFD (0.15),

Revised: ProtAlbert (0.13), ProtXLNet (0.12).

Revised: However, calibration methods such as isotonic regression (Zadrozny & Elkan, 2002) may be applied as a post-processing step to improve the calibration of attention weights and thereby provide a more precise understanding of attention as a measure of conﬁdence.

Revised: We note a limitation of the use of attention as an estimator of contacts: attention is normalized, such that the attention weights originating from a particular token all sum to one, whereas contact maps represent an unnormalized binary scoring function.

Revised: If an amino acid has multiple contacts, therefore, attention will be limited in its ability to provide a well-calibrated estimate.

Revised: In practice, contacts are relatively sparse in the dataset as shown in Figure 21.

Revised: Nonetheless, attention, by virtue of its normalization, will become increasingly diluted as the number of contacts increases

Revised: [Figure]

Revised: r e

Revised: [Figure]

Revised: u e n c

Revised: [Figure]