Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: ### Summary

This paper offers an in-depth analysis of attention in large-scale language models including (AL)BERT and XLNet. Key "headline" findings:

- Attention maps highlight amino acid pairs that are far in sequence space but near in coordinate space.
- Attention maps highlight binding sites within proteins.
- Attention maps capture local secondary structure.
- Attention heads have affinities to specific residue types.

More nuanced findings:

- The evolution across layers of attention-encoded information differs from embedding-encoded information: latter accumulates steadily while former can be isolated to single (end) layer.
- Evidence indicates that attention heads' affinities to specific residue types rests on physical properties rather than mere identity.

### Merits

This is an interesting study that is modest in its characterization of its contributions but diligent in its empirics. Too many papers in this space try to train giant models and cherry-pick tasks for which those models "beat" non-neural baselines, while contributing too little (in the reviewer's humble opinion) to our understanding of what's going on. The authors have shed light on the attention mechanism, showing the intuitive but notable results that attention correlates with coordinate-space distance, binding sites, secondary structure, and residue identity.

I particularly liked the authors' experimental design in 4.4 to shed light on whether identity or structural/functional properties are driving the attention.

The paper is generally well-written and will be accessible to those without prior exposure to protein modeling. The authors' efforts toward accessibility are clear both in the length of background exposition and in the terminology: although I use the term "residue" in this review, I appreciate the authors' choice of "amino acid".

### Drawbacks

**Calibration**: I had to read the "Attention is a well-calibrated predictor of contact maps in some cases" section about 10 times over before I had a grasp on what was going on. Even with my current understanding of the section, I'm extremely skeptical about this specific study. Why do I expect attention (which produces a point on the simplex) to be at all calibrated to a contact map (which is an arbitrary pairwise zero-one matrix)? Suppose we had a very tightly clustered protein and a high-enough radius such that all residues were considered to be in contact. No possible attention map would be well-calibrated to the ground-truth contacts.

I would recommend removal of this section. It is a weak bordering on nonsensical result which seems to have been plucked from a single attention head. I suspect this result is not robust and in any case it does not contribute much conceptual understanding.

**Embeddings vs. attention**: In related work, the paper cites prior work that analyze the output embeddings of large scale models for structural and functional properties. The paper itself focuses on properties captured through attention. As a reader, I'm left somewhat unsatiated. What's captured in embeddings that attention represents less well? What's captured in attention that might not necessarily emerge through a metric space defined in embedding land? The paper missed an opportunity in my opinion to do a systematic comparison of information encoding in attention vs. embeddings, aside from what's shown in Figure 6.

### Recommendation

This paper does not make, nor claim to make, novel methodological contributions in the modeling of proteins. What it does is to study Transformer attention using previously-known techniques but applied to protein sequence models. It overall performs an admirable job of doing so. We at ICLR care not only about "learning representations", but about assessing what precisely those representations have learned, so the paper is certainly appropriate for ICLR.

My main hesitation with recommending acceptance is the whole attention calibration business. The metric is mathematically unsound (contact maps certainly do not sum either row-wise or column-wise to one) and the fact that this paper found a "well-calibrated" attention head isn't all that notable.

With that caveat, I cautiously recommend acceptance. I will raise my score upon either removal of "attention calibration" from the paper or an *extremely* convincing explanation of why I've misinterpreted the metric / why I should care about the calibration of an attention map with an arbitrary unnormalized zero-one matrix.

UPDATE: The authors have removed the attention calibration study. My new score is a 7. This is a good paper and I heartily recommend its acceptance.

### Comments

- p1: "substitution properties" is a handwavy term. Please make the meaning more precise.
- p3: "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)." very good thank you!
- p4: very nice ALBERT visualization and consistent with the model architecture!
- p6: "deeper layers focus relatively more attention on binding sites and contacts". OK, contacts sure, but binding sites?
review_title: Good paper but for "calibration"

final_score: rating: 6: Marginally above acceptance threshold
review_content: The authors of this manuscript conducted a comprehensive analysis of the interpretability of self-attention language models when learning from protein sequences. Specifically, five multi-head self-attention models from NLP were used to model protein sequences. The following tasks were conducted: (1) whether attention captures secondary and tertiary structural information, (2) how and why attention targets binding sites and post-translational modifications, (3) how attention captures properties of varying complexity across different encoding layers, and (4) relationship between pairwise attention amino acids and substitution matrix. Through measures and visualizations, the authors can discover interesting relations between attention and various structural properties. In general, I think it is a good analysis paper. Even though there is no contributions in modelling, the authors can use suitable tools to probe and show the relationships between attention and protein structural informations. I believe the findings of this work is informative and inspiring to researchers in bioinformatics.

This paper can be enhanced if consistency and difference between these five attentions models can be further analyzed in the four tasks. What would be the most suitable model for protein sequences?

Moreover, I would suggest the authors to discuss how many layers and heads are needed for protein sequence modelling. What will happen if 50 or just 5 layers are used for the analysis in this paper? Will the same conclusions be drawn? 

Minors: 
1. transformer -> Transformer
2. The list of references can be further formatted. For example, "Advances in neural information processing systems" -> "Advances in Neural Information Processing Systems".
 

review_title: Systematic analysis of of interpretability of Transformer-based language attention models used for protein sequences

final_score: rating: 7: Good paper, accept
review_content: This is paper is an investigation of the potential biological interpretation of attention in five pretrained protein sequence Transformer models (one from TAPE and four from ProtTrans). It seems that the meatiest contribution in this paper involves the use of the attention analysis metric p_alpha(f) to measure the co-occurrence of attention values (above threshold alpha) and a contact matrix f. A related contribution is the definition of attention calibration, which is a slight variant of Expected Calibration Error applied to evaluating the calibration between attention values and a discretized contact matrix. 

My biggest concern with this paper is that they analyze Transformer models by looking for a single attention head that maximizes p_alpha(f) without, as far as I can tell, really evaluating how often these multiple comparisons might lead to a high p_alpha(f) value by chance. Several compelling examples are shown for individual attention heads (e.g. repeated evaluation of TapeBert 12-4). These individual heads really could just be the locus in which a pretrained model encodes understanding of structural adjacency but...there has to be a null model when trying every node of a big neural network and taking the max value of an evaluation metric. The paper repeatedly plots the background frequency of contacts between residues in the protein data but that's not the same as determining how often your analysis would yield high values just by chance. 

Similarly, what's the null model for the calibration analysis? The plot, like others, has an orange line showing the background frequency of contacts but that seems really unrelated to calibration. I don't know what the null model *should* be but this paper should endeavor to show that finding one calibrated attention head (among the many possibilities) is surprising. 

Other comments:

* In the first background items, on "Amino acids", the description derails into discussion of substitution matrices. If understanding substitution matrices is important to the paper, give them their own bullet point. Otherwise it's a strangely specific technique to discuss while describing amino acids in general. 

* "every protein sequence is formed from a vocabulary of 20 standard amino acids" it's strange to say this and then talk about PTMs a few paragraphs down. If you're presenting PTMs to the networks you're evaluating then you're likely using a larger amino acid alphabet (I think TapeBert has 25 amino acids?)

* This description of secondary structure seems odd: "Secondary structure describes the local segments of proteins, and may be divided into three broad categories: Helix, Strand, and Turn/Fold, as well as an Other category for local structure that falls outside these categories." If you asked most biologists what the broad categories of secondary structure are would probably say "[alpha] Helix" "[beta] sheet", it didn't even occur to me to break sheets into strands and turns. I guess it makes sense that these categories might occur in some secondary structure prediction datasets or tasks, but it seems like the details of some dataset are being presented as representative of biology. 

* Slightly confusing terminology here: "Binding sites are protein regions that bind with other molecules to carry out a specific function. For example, the HIV-1 protease is an enzyme responsible for a critical process in replication of HIV (Brik & Wong, 2003). It has a binding site, shown in Figure 1b, that is a target for drug development to ensure inhibition." I think it's important to distinguish the active sites of a protein where that protein binds some natural target or ligand and the binding site of a drug which inhibits the protein. These often overlap to some degree but the binding site of a small molecule inhibitor can be a much smaller pocket than a natural ligand. 

* "Attention targets PTMs in a small number of heads": I wish it was clearer what this analysis meant. Were the input sequences given as unmodified AAs and then you're evaluating which residues are often observed as modified in vivo? Or did you give the networks sequences with PTMs in them and then this attention pays attention to when rare AAs get used?

OK, after this list of complaints, I want to come back to what I like about this paper. I think many people (myself included) suspect that attention in protein Transformer language models operates in biologically interpretable ways. It's really nice to actually show it and also show that attention is compatible with other simpler biological notions (like AA substitution matrices). But: it's not enough to just look for the best node in each network, you have to make some kind of effort to quantify how like that finding would be by chance in a large enough network. 
review_title: Cool paper but: a little loose with biological concepts and a major statistical concern

final_score: rating: 6: Marginally above acceptance threshold
review_content: General comments:

1) authors claim

"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”. 

and in the related work section

"In our work, we take an interpretability-first perspective to focus on the internal model representations, specifically attention and intermediate hidden states, across multiple protein language models. We also explore novel biological properties including binding sites and post-translational modifications.”

However the structure of evidence is very phenomenological in nature — x (attention at head H, layer L) is observed which coincides with phenomena Y (binding/PTM/). 


2) no cross-model analysis. Why are we observing different outcomes in different models?

Specific comments:

section 4.1 "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”
	- why strongest alignment in only few heads and in high layers?
	- "While there seems to be a strong correlation between the attention head output and classically-defined contacts, there are also differences”. What are the differences?

	- Attention as a well-Calibrated predictor section. Authors offer no explanation why the ECE is higher for none TapeBert. 
		
section 4.2 *why* is the causal reason for attention focusing on binding sites or PTMs?

section 4.4: what about the other 4 amino acids that dont have 25% allocated attention in TapeBert? what is it about those amino acids?

review_title: Good paper, interesting topic/application of attention but explanatory goals is only weakly supported

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

# Major comments
1. The paper is missing two important baselines: 1) a Transformer with randomly initialized weights and 2) a Transformer trained on randomly shuffled protein sequences (with the same amino acid frequencies. These baselines will show if Transformers actually learn protein features, or if the correlation between network features and protein features is an artifact of the Transformer architecture. It has been shown that randomly initialized convolution networks already encode edge detectors in their weights without any training. Something similar can hold true for proteins.

2. To further support that Transformer learns protein features, I would like to see how features specialize during training on protein features. For example, how do figure 2 and figure 6 look at different numbers of training epochs?

3. The attention analysis (equation 1) depends on a threshold \theta. How was it chosen and how sensitive are results to this threshold?

4. Instead of using equation 1 for analyzing ‘alignment’ of attention $\alpha_{i, j}$ with the property $f(i, j)$, I suggest computing the Spearman correlation between  $\alpha_{i, j}$ and  $f(i, j)$ over all samples in the training set. Using the Spearman correlation is a more clear measure for quantifying the relationship between a feature and a property and does not depend on threshold.

5. Using Expected Calibration Error (ECE) for analyzing the calibration of attention is a straight-forward extension of using ECE for analyzing calibration of a ML model, which is widely performed. I also doubt the usefulness of analyzing calibration of attention since practitioners are unlikely to use a particular attention map directly for contact prediction, but train a ML model on top of attention maps to predict contacts (see Rives 2020). What matters is if the resulting ML model is calibrated.

6. For the binding site and post-translational modification (PTM) analysis (section 4.2) it is unclear how $f(i, j)$ (equation 1) is defined. A binding site and (PTM) is a property of a single position $i$ not a pair of positions $(i, j)$.

7. Section 3, datasets: How similar are test sequences to sequences that were used for training Transformer models? Sequences must not overlap and must have a maximum similarity of, e.g., 80%.

8. It is unclear which Transformer architecture was used for Figure 5 and Figure 6. For all analysis, please mention the Transformer architecture and refer to additional results for other architectures. Also discuss in each paragraph if the observation holds true for all architectures or not.

9. Section 4.2, Figure 4: I agree that TapeBert attends most strongly to the last layers but I cannot see a clear pattern for the other architectures (App C3). 

# Minor comments
10. Section 3, Probing task’: Please also cite Biswas 2020 (‘Low-N protein engineering …’).

11. Section 4.1; Figure 2: ‘alignment’ can be misunderstood as ‘sequence alignment’ in the context of protein. Please use another term, e.g. ‘association’ or ‘correlation’.

12. Section 4.1: What are ‘implicit features’?
review_title: Evaluation needs to be strengthened

Meta-review: This paper offers an in-depth analysis of attention in large-scale language models including (AL)BERT and XLNet in the context protein representation learning, and obtains many interesting findings. This is not a typical paper with novel technologies proposed, instead, it studies the existing technologies in a specific (biology) context and explains what the learned representations and attention map really mean. 
All the reviewers see the value in this paper and give positive feedback in general. At the same time, they also raised a few concerns, e.g., regarding the claim on “well-calibrated" attention head, on some missing details of the algorithm description and the experiments,  on phenomenon vs. causality of the finding, etc. The authors really did a very good job in their rebuttal and paper revision, and most of these concerns were (at least partially) addressed, and a few reviewers raised their scores. With this, we are quite confident that this paper is above the bar of ICLR. 

Decision: Accept (Poster)