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 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.
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 how attention is consistent with a classic measure of similarity between amino acids—the substitution matrix.
Revised: Further, we show how attention is consistent with a classic measure of similarity between amino acids—the substitution matrix.

Original: Finally, we demonstrate that attention captures progressively higher-level representations of structure and function with increasing layer depth.
Revised: Finally, we demonstrate that attention captures progressively higher-level representations of structure and function with increasing layer depth.

Original: [Figure]
Revised: [Figure]

Original: 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 of amino acids, of which 20 are commonly observed.
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.

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

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: 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; two commonly observed types are the alpha helix and beta sheet .
Revised: Secondary structure describes the local segments of proteins; two commonly observed types are the alpha helix and beta sheet .

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 (proteins, natural ligands, and small-molecule drugs) 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 HIV-1 protease 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 (El-Gebali 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
Revised: We also analyze 4 pre-trained

Original: 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
Revised: 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: We also present an alternative, continuous version of this metric in Appendix B.1.
Revised: We also present an alternative, continuous version of this metric in Appendix B.1.

Original: 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).
Revised: 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 the property (e.g. the proportion of attention focused on binding sites).
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).

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

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

Original: Bonferroni correction for multiple hypotheses corresponding to multiple attention heads.
Revised: Bonferroni correction for multiple hypotheses corresponding to multiple attention heads.

Original: Second, we compare the results to a null model, which is an instance of the model with randomly shufﬂed attention weights.
Revised: Second, we compare the results to a null model, which is an instance of the model with randomly shufﬂed attention weights.

Original: We describe these methods in detail in Appendix B.2.
Revised: 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 (Reif 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.3.
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.

Original: We considered a more ﬁne-grained taxonomy of secondary structure with three categories: Helix ,
Revised: We considered a more ﬁne-grained taxonomy of secondary structure with three categories: Helix ,

Original: Strand , and Turn/Bend , with the last two belonging to the higher-level beta sheet category (Sec. 2).
Revised: Strand , and Turn/Bend , with the last two belonging to the higher-level beta sheet category (Sec. 2).

Original: We used this taxonomy in to study how the model understood structurally distinct regions of beta sheets.
Revised: 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
Revised: We obtained token-level binding site and protein modiﬁcation labels from the Protein Data

Original: Bank (Berman et al., 2000).
Revised: 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
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

Original: [Figure]
Revised: [Figure]

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

Original: We only include results for attention heads where at least 100 high-conﬁdence attention arcs are available for analysis.
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.
Revised: We set the attention threshold θ to 0.

Original: to select for high-conﬁdence attention while retaining sufﬁcient data for analysis.
Revised: 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).
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 how attention aligns with contact maps across the heads of the ﬁve models evaluated 2 , based on the metric deﬁned in
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

Original: Equation 1.
Revised: Equation 1.

Original: 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%.
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 the top head in TapeBert.
Revised: Figure 1a shows an example protein and the induced attention from the top head in TapeBert.

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

Original: [Figure]
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: 4. B INDING S ITES AND P OST -T RANSLATIONAL M ODIFICATIONS
Revised: 4. B INDING S ITES AND P OST -T RANSLATIONAL M ODIFICATIONS

Original: We also analyze how attention interacts with binding sites and post-translational modiﬁcations (PTMs), which both play a key role in protein function.
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: Attention targets binding sites throughout most layers of the models.
Revised: Attention targets binding sites throughout most layers of the models.

Original: Figure 3 shows the proportion of attention focused on binding sites (Eq. 1) across the heads of the 5 models studied.
Revised: Figure 3 shows the proportion of attention focused on binding sites (Eq. 1) across the heads of the 5 models studied.

Original: 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%.
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: 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.
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: Figure 1b visualizes the attention in one strongly-aligned head from the TapeBert model.
Revised: Figure 1b visualizes the attention in one strongly-aligned head from the TapeBert model.

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

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

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

Original: Appendix A.1), whereas the other 4 models all use masked language modeling objectives.
Revised: Appendix A.1), whereas the other 4 models all use masked language modeling objectives.

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

Original: The embedding probes (orange) quantify the knowledge of the given property that is encoded in each layer’s output embeddings.
Revised: The embedding probes (orange) quantify the knowledge of the given property that is encoded in each layer’s output embeddings.

Original: The attention probe (blue), show the amount of information encoded in attention weights for the (pairwise) contact feature.
Revised: The attention probe (blue), show the amount of information encoded in attention weights for the (pairwise) contact feature.

Original: Additional details are provided in Appendix B.3.
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.

Original: By attending to these regions, the model can leverage this higher-level context when predicting masked tokens throughout the sequence.
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 amino acids associated with 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 PTM positions, though these occur at only 0.8% of sequence positions in the dataset.
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: 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.
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
Revised: We analyze how attention captures properties of varying complexity across different layers of

Original: TapeBert, and compare to a probing analysis of embeddings and attention weights (see Section 3).
Revised: 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 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.
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.

Original: The probing analysis
Revised: The probing analysis

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: QRSTVWY
Revised: QRSTVWY

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

Original: of attention
Revised: of attention

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

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 6 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%.
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: Similar behavior was observed for ProtBert, ProtBert-BFD,
Revised: Similar behavior was observed for ProtBert, ProtBert-BFD,

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

Original: Detailed results for TapeBert including statistical signiﬁcance tests and comparison to a null model are presented in Appendix C.5.
Revised: 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 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.
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
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

Original: (left) for TapeBert.
Revised: (left) for TapeBert.

Original: For example, the entry for Pro (P) and Phe (F) is the correlation between the two heatmaps in Figure 6.
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 7 (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 the ProtTrans models: 0.68 (ProtBert), 0.75 (ProtBert-BFD), 0.60 (ProtAlbert), and 0.71 (ProtXLNet).
Revised: Similar correlations are observed for the ProtTrans models: 0.68 (ProtBert), 0.75 (ProtBert-BFD), 0.60 (ProtAlbert), and 0.71 (ProtXLNet).

Original: As a baseline, the randomized versions of these models (Appendix B.2) yielded correlations of -0.02 (TapeBert), 0.
Revised: As a baseline, the randomized versions of these models (Appendix B.2) yielded correlations of -0.02 (TapeBert), 0.

Original: (ProtBert), -0.03 (ProtBert-BFD), -0.05 (ProtAlbert), and 0.21 (ProtXLNet).
Revised: (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.

Original: 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).
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 by considering particular biophysical properties and relationships.
Revised: Our work extends these methods to protein sequence models by considering particular biophysical properties and relationships.

Original: We also present a joint cross-layer probing analysis of 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 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.
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.

Original: Finally, we present a novel tool for visualizing attention embedded in three-dimensional structure.
Revised: 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: 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 AGREEMENT 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 matches 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. S TATISTICAL SIGNIFICANCE TESTING AND NULL MODELS
Revised: B. S TATISTICAL SIGNIFICANCE TESTING AND NULL MODELS

Original: We perform statistical signiﬁcance tests to determine whether any results based on the metric deﬁned in Equation 1 are due to chance.
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.

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

Original: ( i, j )
Revised: ( i, j )

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

Original: ( i, j )
Revised: ( i, j )

Original: = 1 .
Revised: = 1 .

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

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

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

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

Original: We initially considered using randomly initialized versions ofthe models studied as null models.
Revised: We initially considered using two forms of null models: (1) a model with randomly initialized weights. and (2) a model trained on randomly shufﬂed sequences.

Original: However,  none of these models yielded attention weights greater than the attention threshold θ .
Revised: However, in both cases, none of the sequences in the dataset yielded attention weights greater than the attention threshold θ .

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

Original: 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: Therefore, we implemented an alternative randomization scheme in which we randomly shufﬂe attention weights from the original models as a post-processing step.

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

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

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

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

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

Original: B. P ROBING METHODOLOGY
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 4 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 do 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

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table] [Table]
Revised: [Table] [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]

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

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: A tt e n t i o n %
Revised: A tt e n t i o n %

Original: [Figure]
Revised: [Figure]

Original:  i o
Revised: t i o

Original: [Figure]
Revised: [Figure]

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

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

Original: t i o
Revised: t i o

Original: [Figure]
Revised: [Figure]

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

Original: [Figure]
Revised: [Figure]

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

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

Original: [Table]
Revised: [Table]

Original: Arg
Revised: Arg

Original: R
Revised: R

Original: Arginine 5.
Revised: Arginine 5.

Original: Asn
Revised: Asn

Original: N
Revised: N

Original: Asparagine 4.
Revised: Asparagine 4.

Original: Asp
Revised: Asp

Original: D
Revised: D

Original: Aspartic acid 5.
Revised: Aspartic acid 5.

Original: Cys
Revised: Cys

Original: C
Revised: C

Original: Cysteine 1.
Revised: Cysteine 1.

Original: Gln
Revised: Gln

Original: Q
Revised: Q

Original: Glutamine 3.
Revised: Glutamine 3.

Original: Glu
Revised: Glu

Original: E
Revised: E

Original: Glutamic acid 6.
Revised: Glutamic acid 6.

Original: Gly
Revised: Gly

Original: G
Revised: G

Original: Glycine 7.
Revised: Glycine 7.

Original: His
Revised: His

Original: H
Revised: H

Original: Histidine 2.
Revised: Histidine 2.

Original: Ile
Revised: Ile

Original: I
Revised: I

Original: Isoleucine 5.
Revised: Isoleucine 5.

Original: Leu
Revised: Leu

Original: L
Revised: L

Original: Leucine 9.
Revised: Leucine 9.

Original: Lys
Revised: Lys

Original: K
Revised: K

Original: Lysine 6.
Revised: Lysine 6.

Original: Met
Revised: Met

Original: M
Revised: M

Original: Methionine 2.
Revised: Methionine 2.

Original: Phe
Revised: Phe

Original: F
Revised: F

Original: Phenylalanine 3.
Revised: Phenylalanine 3.

Original: Pro
Revised: Pro

Original: P
Revised: P

Original: Proline 4.
Revised: Proline 4.

Original: Ser
Revised: Ser

Original: S
Revised: S

Original: Serine 6.
Revised: Serine 6.

Original: Thr
Revised: Thr

Original: T
Revised: T

Original: Threonine 5.
Revised: Threonine 5.

Original: Trp
Revised: Trp

Original: W
Revised: W

Original: Tryptophan 1.
Revised: Tryptophan 1.

Original: Tyr
Revised: Tyr

Original: Y
Revised: Y

Original: Tyrosine 3.
Revised: Tyrosine 3.

Original: [Table]
Revised: [Table]

Original: D A TTENTION C ALIBRATION A NALYSIS

Original: [Figure]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

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

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

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.

Original: We demonstrate this approach by measuring attention calibration with the contact map property f

Original: ( i, j ) , a pairwise feature indicating whether the amino acids at positions i and j are in contact.

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 plots in Figure 20 ( x -axis).

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

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

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

Original: 1 ) would dominate this calculation.

Original: The ﬁrst bin in Figure 20a, for example, contains 99.96% of all examples.

Original: Since both attention and empirical probability is low in this bin, a weighted average would produce an artiﬁcially low calibration error.

Original: Therefore we instead use an unweighted average.

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

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

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

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

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

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

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

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

Original: [Figure]

Original: r e

Original: [Figure]

Original: u e n c

Original: [Figure]