Single Layers of Attention Suffice to Predict Protein Contacts
Authors: Nick Bhattacharya
Neil Thomas
Roshan Rao
Justas Daupras
Peter K Koo
David Baker
Yun S. Song
Sergey Ovchinnikov
Venue:  (2021)
Keywords: Protein Structure
Proteins
Contact Prediction
Representation Learning
Language Modeling
Attention
Transformer
BERT
Markov Random Fields
Potts Models
Self-supervised learning

Abstract: The established approach to unsupervised protein contact prediction estimates coevolving positions using undirected graphical models. This approach trains a Potts model on a Multiple Sequence Alignment, then predicts that the edges with highest weight correspond to contacts in the 3D structure. On the other hand, increasingly large Transformers are being pretrained on protein sequence databases but have demonstrated mixed results for downstream tasks, including contact prediction. This has sparked discussion about the role of scale and attention-based models in unsupervised protein representation learning. We argue that attention is a principled model of protein interactions, grounded in real properties of protein family data. We introduce a simplified attention layer, factored attention, and show that it achieves comparable performance to Potts models, while sharing parameters both within and across families. Further, we extract contacts from the attention maps of a pretrained Transformer and show they perform competitively with the other two approaches. This provides evidence that large-scale pretraining can learn meaningful protein features when presented with unlabeled and unaligned data. We contrast factored attention with the Transformer to indicate that the Transformer leverages hierarchical signal in protein family databases not captured by our single-layer models. This raises the exciting possibility for the development of powerful structured models of protein family databases.

Under review as a conference paper at ICLR 2021
SINGLE LAYERS OF ATTENTION SUFFICE TO PREDICT PROTEIN CONTACTS
Anonymous authors Paper under double-blind review
ABSTRACT
The established approach to unsupervised protein contact prediction estimates co- evolving positions using undirected graphical models. This approach trains a Potts model on Multiple Sequence Alignment, then predicts that the edges with high- est weight correspond to contacts in the 3D structure. On the other hand, in- creasingly large Transformers are being pretrained protein sequence databases but have demonstrated mixed results for downstream tasks, including contact pre- diction. This has sparked discussion about the role of scale and attention- beaeed models in unsupervised protein representation learning. We argue that attention is principled model of protein interactions, grounded in real properties of pro- tein family data. We introduce a simplified attention layer, factored attention, and show that it achieves comparable performance to Potts models, while shar- ing parameters both within and across families. Further, we extract contacts from the attention maps of a pretrained Transformer and show they perform compet- itively with the other two approaches. This provides evidence that large-scale pretraining can learn meaningful protein features when presented with unlabeled and unaligned data. We contrast factored attention with the Transformer to indi- cate that the Transformer leverages hierarchical signal in protein family databases not captured by our single-layer models. This raises the exciting possibility for the development powerful structured models of protein family databases.
INTRODUCTION
Inferring protein structure from sequence a longstanding problem in computational biochemistry. Potts models, particular kind of Markov Random Field (MRF), are the predominant unsupervised method for modeling interactions between amino acids. Potts models are trained to maximize pseu- dolikelihood alignments of evolutionarily related proteins (Balakrishnan et al., 201T Ekeberg jt al. 2013; Seemayer et al. 2014). Features derived from Potts models were the main drivers of performance at the CASPIT competition (Monastyrskyy al. (016) and have become standard in
state-of-the-art supervised models (Wang et al.) 2017 Yang et 2019 Senior et al. 2020).
Inspired by the success of BERT (Devlin et al. 2018). GPT (Brown et al. 2020) and related un- supervised models in NLP, line of work has emerged that learns features of proteins through
self-supervised pretraining (Rives et al.) [2020] Elnaggar et al. 2020 Rao et al. 22019) Madani al.)
2020 Nambiar et al.) 2020). This new approach trains Transformer (Vaswani et al. 2017) models on large datasets of protein sequences. There is significant debate over the role of pretraining in protein modeling. Pretrained model performance raises questions about the importance of data and model scale (Lu et al.] 2020 Elnaggar et al. 2020), the potential for neural features to compete with evolutionary features extracted by established bioinformatic methods (Rao et al.) 2019). and the benefits of transfer learning for protein landscape prediction (Shanehsazzadeh et al. 2020).
In this paper, we take the position that attention-basse models can build on the strengths of both Potts trained on alignments and Transformers pretrained on databases. We introduce a simplified model, factored attention, and show that it is motivated directly by fundamental properties of protein sequences. We demonstrate empirically that a single layer of factored attention suffices to predict protein contacts competitively with state-of-the--e -fr Potts models, while leveraging parameter sharing across positions within a single family and across hundreds of families. This highlights the potential for explicit modeling of biological properties with attention mechanisms.

---

Under review as a conference paper at ICLR 2021
Further, we systematically demonstrate that contacts extracted from ProtBERT-BFU Elnaggar et al. 2020) are competitive with those estimated by Potts models across 748 protein families, inspired by recent work from Vig al. (2020). We find that ProBBERT-BFD outperforms Potts for proteins with alignments smaller than 256 sequences. This indicates that large-scale Transformer pretraining merits continued efforts from the community.
We contrast factored attention with ProtBERT-BFD to identify a large gap between the gains af- forded by the assumptions of single-layer models and the performance of Pro(BERT-BFD. This suggests the existence of properties linking protein families that allow for effective modeling of thousands of families at once. These properties are inaccessible Potts models, yet are implicitly learned by Transformers through extensive pretraining. Understanding and leveraging these proper- ties represents an exciting challenge for protein representation learning.
Our contributions are as follows:
1. We analyze the assumptions made by the attention mechanism in the Transformer and show that a single attention layer a well-founded model interactions within protein families; attention for proteins can be justified without any analogies to natural language.
2. We show that single-layer models can successfully share parameters across positions within family or sharing of amino acid features across hundreds of families, demonstrating that factored attention achieves performance nearly identical to Potts with far fewer parameters.
3. We carefully benchmark ProtBERT-BED against an optimized Potts implementation and show that the pretrained Transformer extracts contacts competitively with Potts.
2 BACKGROUND
Proteins are polymers composed of amino acids and are commonly represented as strings. Along with this ID sequence representation, each protein folds into a 3D physical structure. Physical distance between positions in 3D is often a much better indicator of functional interaction than proximity sequence. One representation of physical distance is a contact map C, a symmetric matrix in which entry Cj = if the 3 carbons of i and are within 8Å of one another, and O otherwise.
Multiple Sequence A Alignments. To understand structure and function of protein sequence, one typically assembles set of its evolutionary relatives and looks for patterns within the set. A set of related sequences is referred to as a protein family, commonly represented by Multiple Sequence Alignment (MSA). Gaps in aligned sequences correspond to insertions from an alignment algorithm JJhnsson al., 2010; Remmert et al. 2012), ensuring that positions with similar structure and function line up for all members of the family. After aligning, sequence position carries significant evolutionary, structural, and functional information. See Appendix/A.I]foo more information.
Cuevolutionary Analysis of Protein Families. The observation that statistical patterns in MSAs be used to predict couplings has been widely used infer structure and function from protein
families (Korber et al. 1993] Göbel et af.) 1994 Lapedes et al.) 1999 Lockless & Ranganathan
1999 Fodor& & Aldrich 2004 Thomas al. 2008 Weigt et al.) 2209 Fatakia al. [2009]. Let X;
be the amino acid at position i sampled from a particular family. High mutual information between X, and Xj suggests an interaction between positions and j. The main challenge in estimating contacts from mutual information is to disentangle "direct couplings" corresponding to functional interactions from interactions induced by con-functional patterns (Lapedes al. 1999 Weigt al., 2009). State-of-the-ar estimates interactions from MRF parameters, as described below.
Supervised Structure Prediction. Modern structure prediction methods take a supervised ap- proach, taking MSAs as inputs and outputting predicted structural features. Deep Learning has greatly advanced state of the art for supervised contact prediction (Wang et al. [017 Jones & Kan- dathil 2018 Senior et al.) 2020 Adhikari 2020). These methods train deep residual networks that take covariance statistics or cocvolutionary parameters as inputs and output contact maps or distance matrices. Extraction of contacts without supervised structural signal has not competitive per- formance from neural networks until the recent introduction of Transformers pretrained on protein databases.
2

---

Under review as a conference paper at ICLR 2021
W41
Figure 1: Training and interpretation of Potts model on single protein family. The proteins are all loops formed by one blue and one yellow amino acid locking together. The MSA for this family aligns these critical yellow and blue amino acids. For the trained MRF on this MSA, the weight matrix 1 has the highest values due to the evolutionary constraint that blue and yellow covary for those positions. this case, the highest predict contact recapitulates true contact.
Attention-Based Protein Models. Elnaggar al. (2020); Rives et al. (2020): Rao al. (2019) ;
Madani et al. (2020) train unsupervised attention- models on millions to billions of protein sequences and show that pretraining can produce good representations. Rao et al. (2019) and Rives et al. (2020) assess embeddings
inputs to supervised structure prediction methods. When protein structure known, Ingraham et al. (2019) and Du et al. (2019) have used sparse attention based on the physical structure of the protein to directly encode priors about interaction.
METHODS
Throughout this section, x (x1, (..,..., ,., is a sequence of length L from an alphabet of size A. This sequence is part of MSA of length L with N total sequences. Recall that fully-connected Pairwise MRF over p variables Xi...., ...,,,, Xp specifies a distribution
(1 , Fp) Z exp
Σ Fo (o)
where is the partition function and Eo(x; 22)) is an arbitrary function of i, J. x, and øy. For all models below, introduce explicit functional Fo((:;) capture the marginal distribution of X,. When introduced, we parametrize the marginal with Ex(x;) = bi for E RLX
3.1 POTTS MODELS
A Potts model is fully-connected pairwise MRF with L variables, each representing a position in the MSA. An edge (i.j) parametrized with matrix Wij AxA These matrices are organized into an order-4 tensor which form the parameters of Potts model, see Figure 1 Note that Wr5 Wi. The energy functional Potts model is given through lookups, namely
Eo 2Ty) = W (1),))z).
3.2 FACTORED ATTENTION
Like Potts, factored attention is fully-connected pairwise MRF with L variables. The parameters of this model consist of H triples (Wq. Wx. Wv). where Wo, Wx E RLxd, Wy E Ax and d a hyperparameter. Each such triple called a head and d the head size. Unlike Potts model, the parameters for each edge (i,j) are tied through the use of heads. The energy functional is
Eg(,,, IT))
Σ
symm (softmax (wb (ūn
(2)
h=1
where symm(M) (M M2 ensures the positional interactions symmetric.
3

---

Under review as a conference paper at ICLR 2021
Factored attention has two advantages over Potts for modeling protein families: it shares pool of amino acid feature matrices across all positions and it estimates (L) parameters instead of (
Sharing amino acid features. Many contacts in a protein are driven by similar interactions between amino acids, such as many types of weakly polar interactions (Burley & Petsko 1988 Daenicke 2000). If two pairs of positions (i and (1, m) both in contact due to the same interaction, Potts model must estimate completely separate amino acid features Wo and Wit. In order to share amino acid features, we want to compute all energies from one pool of A x A feature matrices. The simplest way to accomplish this is by associating an L x L matrix A to every A x A feature matrix Wy. For H such pairs (A, Wv), we could introduce a factorized MRF:
E,( symm (softmax (4 (x),,)))
(3)
A row-wise softmax is taken to encourage sparse interactions and aid in normalization. This model allows the pairs (i.j) and (1, m) to reuse single feature WV,, assuming Aly and Al, are both large.
Scaling linearly in length. Both Potts and the factorized model in Equation 3 have O(L2) param- eters. However, contacts observed to grow linearly over the wide range of protein structures currently available (Taylor & Sadowski, [001] Kamisetty et al.] 2013). which we examine in Figure I1 Given that the number of interactions we wish estimate grows linearly in length, the quadratic scaling of these models can be greatly improved. One way to fix this is by introducing the factoriza- tion A WoWk k. where Wq. WK E Rixdd and normalization.
As before, we employ a row-wise softmax for sparsity
Combining feature sharing with linear length scaling leads to Equation 2
Recovering attention with sequence-dependent interactions. All models introduced so far esti- mate a single undirected graphical model from the training data. While a single graph can be good approximation for the structure associated with a protein family, many families have subfamilies with different functional specializations and even different underlying contacts (Brown et al. [000)] Malinverni & Barducci| 2019)). Since subfamily identity is rarely known, allowing edge weights to be a function of sequence could enable the estimation of family of highly related graphs
The attention mechanism of the Transformer implements sequence-dependent edge weights by al- lowing positional interactions to be a function A(x). In the language of the Transformer, factored attention estimates single graph because it computes queries and keys using only the positional encoding. More precise explanations are in Section A.2
While data-dependent interactions significantly increase model capacity, this comes with multiple tradeoffs. Demands on GPU memory increase significantly with data-dependent interactions, as detail in Section
A.3 Further, attention- based models can be very challenging to train (Zhang jet al. 2019 Liu et al. 2019); we found both Potts and factored attention trained more stably in all experiments. Lastly, the expressivity conferred by attention creates risk of overfitting, especially when the pretraining dataset is small. Effective training on small MSAs is particularly important, since MSA depth correlates strongly with contact quality (Schaarschmidt et al. 2018).
Self-Supervised Losses, Given an MSA, tfute-of-thheart methods estimate Potts model parameters through pseudolikelihood maximization (Kamisetty et al.) [2013] Ekeberg et al. 2013). On the other hand, BERT-like attention models are typically trained with variants of masked language modeling (Devlin al.) 2018). Pseudolikelihood is challenging to compute efficiently for generic models, unlike masked language modeling. Both of these losses require computing conditionals of the form po (rs/ where M subset of [I....;L} L} \ (i). The losses LPt. and LMIMM for pseudolikelihood and masked language modeling, respectively, are
(f;(0; =
log logro(z)l
(0),
Σ log (o iEM
Regularization for Potts and factored attention are both based on MRF edge parameters, while single-layer attention is penalized using weight decay. More details can be found in Section
4

---

Under review as a conference paper at ICLR 2021
P@L 0.68
P@L 0.71
p@L 0.50
P@L L 0.78
50
100
150
Potts
Factored Attn
Single-Layer Attn
ProtBERT-BE
Figure 2: Predicted contact maps and Precision at for each model on PDB entry 2BFW. Blue indicates true positive, red indicates a false positive, and grey indicates a false negative.
3.3 EXTRACTING CONTACTS
Potts. We follow standard practice and extract contact map Ō R4XL from the order-4 interaction tensor W by setting Cij ||W*)|| ||..
Factored Attention. Since factored attention is a pairwise MRF, we can compute its order-4 inter- action tensor W and use the same procedure as Potts. See Equation
Single-Layer Attention. To produce contacts for an MSA, we compute attention maps from only the positional encoding (without sequence) and average attention maps from all heads. Each single- layer attention model is trained on one MSA, so the positional encoding is a feature shared by all sequences in the MSA.
ProtBERT-BFD. We extract contacts from ProtBERT by averaging a subset attention maps for an input sequence x. Of the 16 heads in 30 layers, selected six whose attention maps had the top individual contact precisions over 500 families randomly selected from the Yang et al. (2019) dataset. Predicted contacts for x are given by averaging the L attention maps from these heads, then symmetrizing additively. See Table
Average Product Correction (APC). Empirically, Potts models trained with Frobenius norm regu- larization have artifacts in the outputs O. These are removed with the Average Product Correction (APC) (Dunn et al. 2008 Unless otherwise stated, we apply APC to all extracted contacts.
4 RESULTS
Experimental Setup. We use set of 748 protein families from Yang et al. (2019) to evaluate all models. For Pouts models and single attention layers, we train separate models on each individ- ual MSA. ProBBERT-BFD is frozen for all experiments. We train models using PyTorchL.ightnin
Falcon 2 and Weights and Biases (Biewald, 2020). We compare predicted contact maps Ć to true contact maps C using standard metrics based on precision. A particularly important metric is precision L, where L is the length of the sequence (Schaarschmidt al. 2018 Shrestha al.) 2019). This is computed by masking Ć only consider positions 2 apart, predicting the top entries to be contacts, and computing precision. We provide more information on data and metrics Appendix A.6 and more information on model hyperparameters in Section A.7
Attention assumptions reflected in 15,051 protein structures. We examine all 15,051 structures in the Yang et al. (2019) dataset for evidence of two key properties useful for single-layer attention models: few contacts per residue and the number of contacts scaling linearly in length. Figure []] shows a linear trend of number contacts versus length. In Figure [2] we see that 80% of the 3,747,101 million residues in these structures have 4 or fewer contacts. Only 1.8% of residues have more than ten contacts. This shows that the row-wise softmax, which encourages each residue to attend to only few other residues per-head, reflects structure found in the data.
Factored attention matches Potts performance on 748 families. Figure 2 shows a representative sample of good quality contact maps extracted from all models. Figure summarizes the perfor-

---

Under review as a conference paper at ICLR 2021
1.0
1.0
Potts
0.8
Factored Atta =H 256)
0.8
Precision at
Procision
0.6
Single-Layer
ProBBEKT-BFD
0.4
0.4
0.2
0.2
0.0
0.0
256
1024
4096
16384
200
600
800
(a) MSA depth
(b) Reference length
Figure 3: Model performance evaluated MSA depth and reference length.
1.0
L.
1.0
Predison L/5
0.8
Predison US
0.8
Lis
0.6
0.4
0.2
0.2
0.0
0.0
0.0
Short-range
Medium-range
(c) Long-ran range
Figure 4: Contact precision for all models stratified by the range of the interaction.
mance of all models over the set of 748 protein families. Factored attention, Potts, and ProtBERT- BFD have comparable overall performance, with median precision L of 0.46, 0.47, and 0.48, spectively. Stratifying by number of sequences reveals that ProtBERT-BFI has higher precision on MSAs with fewer than 256 sequences. For MSAs with greater than 1024 sequences, Potts, factored attention, and ProtBERT-BFD have comparable performance. Single-layer attention is uniformly worse over all MSA depths.
Next, we evaluate the impact of sequence length performance. Figure 3b shows that factored attention and Potts achieve similar precision at L over the whole range of family lengths, despite factored attention having far fewer parameters for long families. This shows that factored attention can successfully leverage sparsity assumptions where they are most useful.
Long-range contacts particularly important for downstream structure preduccionn algorithms long-range precision at L/5 reported both CASP12 and CASP13 (Schaarschmidt et al. 2018 Shrestha et al., 2019). Figure 4 breaks down contact precisions based on position separation into short (6 < sep < 12), medium (12 < sep < 24), and long (24 < sep). We see that ProtBERT- BFD performs best on short-range contacts, with median increase of 0.068 precision at L/5. On long-range ProBBERT-BED, there is no appreciable difference in performance to Potts and factored attention. Across the range of contact bins, factored attention and Potts perform very similarly.
Fewer heads can match Potts on L/5 contacts. We probe the limits of parameter sharing by lowering the number of heads in factored attention and evaluating whether fewer heads can be used to precisely estimate contacts. Figure Salshows that 128 heads can be used to estimate L/5 contacts as precisely as Potts over the full set of 748 families. In Figure 5b) we see that factored attention with 32 and 64 heads is still able to achieve reasonable overall performance compared to Potts. 32 and heads have precision at L/5 at least as high as Potts for 329 and 348 families, respectively. wish to recover the top L contacts, 256 heads are required match Potts across all families, as seen in Figure L3 Having more heads than 256 does not further increase performance. Intriguingly, Figure 14 demonstratess that both Spearman and Pearson correlation between the order-4 interaction tensors of factored attention and Potts improve even when increasing to 512 heads. We do not observe the same trends for increasing head size, shown in Figure]

---

Under review as a conference paper at ICLR 2021
1.00
Factored Atto (If= 128)
16384
of MSA sequenies
Precision L.5
0.75
0.75
4096
0.50
0.50
1024
0.25
0.25
256
0.00
0.00
0.0
0.5
1.0
8
16
32
64
128
256
512 Potts
Potts
Number of heads
(a) heads versus Potts P@ L/5
(b) P@ L/5 varying number heads.
Figure 5: Examining impact of number of heads precision L/5. Left: Comparing performance of Potts and 128 heads over each family shows comparable performance. Right: Precision at L/5 drops off slowly until 32 heads, then steeply declines beyond that.
P@ L/5 0.78
P@ L/5 0.86
Precision L/5
0.50
200
Potts
0.25
256
H
0.00
400
10'
10
10
10*
Polts
Training Step
(a) Top L/5 contacts 3n2a.
(b) Training dynamics models on 3n2a.
Figure 6: Factored attention with 4 heads can learn the top L/5 contacts on 3n2a.
For some families, the number of heads can be reduced even further. We show an example on the MSA built for PDB entry 3n2a. In Figure we see that merely 4 heads are required to recover 1./5 contacts nearly identical to those recovered by Potts. This shows that shared amino acid features and interaction parameters can enable identical performance with 300x reduction in parameters. The training dynamics of these models are shown in Figure 6b Both factored attention with 256 heads and Potts converge after roughly 100 gradient steps, whereas factored attention with heads requires nearly 10,000 steps to converge. In Fig we show that the top L contacts significantly worse for 4 heads compared to Potts.
One set of amino acid features be used for all fam- ilies. Thus far we have only examined models that share parameters within single protein families. Since Prot- BERT is trained on an entire database, it can leverage feature sharing across families attain greater parame- ter efficiency and improved performance on small MSAs.
Factored Astn Froozn Values
1.00
16384
0.75
1096
0.50
1024
MSA
0.25
To explore the possibility that attention can share param- eters across families, we train factored attention using a single set of frozen value matrices. We first train factored attention normally on 3n2a with 256 heads, then freeze the learned value matrices for the remaining 747 fami- lies. The query and key parameters are trained normally. In Figure 7 we compare the precision at of factored attention with frozen 3n2a features to that of factored at- tention trained normally. Using a single frozen set of fea- tures results in only 6 families seeing precision at I. decrease by than 0.05, with maximum
256
0.00
0.0
0.5
1.0
Factored Artn
Figure 7: A single set of frozen value matrices can be used for all families.

---

Under review as a conference paper at ICLR 2021
drop of 0.11. Frozen values are also comparable to Potts performance, expected see Figure 17 This suggests that, even for a single-layer model, a single set of value matrices can capture amino acid features actoss functionally and structurally distinct protein families.
Factored attention reduces total parameters estimated. For an MSA of length L with alphabet size A, Potts models require 42) parameters. Factored attention with H heads and head size d requires (2Ldd + A²) parameters. In Figure [8] we plot number of parameters versus length for various values of H and d 32. Potts requires total of 12 billion parameters to model all 748 families. Factored attention with 256 heads and head size 32 has 3.2 billion parameters, lowering 128 heads reduces this 790 million. Half of this reduction comes from 107 families of length greater than 400. ProtBERT-BFD is the most efficient, with 420 million parameters.
Ablations APC has a considerable impact on both Potts and factored attention, creating a median increase in precision at L of 0.1 and 0.07, respectively. The effect of APC is negligible for single- layer attention and ProtBERT. Replacing pseudolikelihood maximization with Masked Language Modeling did not appreciably change performance for either Potts or factored attention. Addition of the single-site potential D' increases performance slightly for attention layers, but not enough to change overall trends. To compare to ProIBERT-BFD. we train our single-layer attention models on unaligned families and found that performance degrades significantly, See Figures 22
5 DISCUSSION
We have shown that attention-based models achieve state-of-the-art performance on unsupervised contact extraction and that the assumptions encoded by attention reflect important properties of protein families. These results suggest that attention has a natural role in protein representation learning, without analogy to attention's 's success the domain of NL..
The motivating principles of factored attention were 1) sharing amino acid features across positions and 2) reducing parameters to OL)). The additional assumption of sequence-dependent contacts led to attention. The parametrizations of these properties presented in this paper were chosen to match those of the Transformer, but many alternatives exist. We believe broader exploration of this space will be fruitful. For example, the importance of APC for factored attention suggests including it directly into the attention mechanism, rather than as a correction after training. The success of Potts and factored attention demonstrates that estimating single graphical model for single protein family is often a reasonable approximation. Adapting attention to better model a set of highly related graphs could prove useful for capturing subfamily structure within deep MSAs, or for transferring information across related families. There remains, unexplored, a rich space of protein- specific architectures that potentially have the capacity to learn powerful protein features.
Many mysteries surround the success of Transformers at extracting protein contacts through masked language modeling. Most surprising is the apparent ability of ProtBERT model many protein families at once, despite not being presented any information about protein family identity or even the existence of such a classification. Beyond family information, ProtBERT appears to learn parsi- monious representations which share a vast amount of signal across protein families. Even including a number of within-famil hierarchical assumptions and sharing amino acid features across all fam- ilies, factored attention falls far short of ProtBERT in its potential to efficiently represent tens of thousands of protein families. Despite its impressive performance, ProtBERT does not improve long-range contact extraction in our evaluations. Much work remains to understand how pretraining can be most useful for downstream applications.
Our results show that hierarchical structure, both within and across families, is source of signal available to attention models. Understanding how such structure can be learned without the use of protein family labels could lead the development of widely applicable modeling components for protein representation learning. We believe that unsupervised attention-base models have the potential to impact structure prediction broadly as existing coevolutionary methods.
REFERENCES
Badri Adhikari. Deepcon: protein contact prediction using dilated convolutional neural networks with dropout. Bioinformatics, 36(2):470-477, 2020.
8

---

Under review as a conference paper at ICLR 2021
Sivaraman Balakrishnan, Hetunandan Kamisetty, Jaime G. Carbonell, Su-In Lee, and Christo- pher James Langmead. Learning generative models for protein fold families. Proteins:
Structure. Function, and Bioinformatics. 19(4):1061-1078) apr 2011. ISSN 08873585. doi:
10.1002/prot.229334 URL Ut.httt://ddii wi ey. cwm/co..c..cmmmo /tt.//0/.../////2.
Lukas Biewald. Experiment tracking with weights and biases, 2020.
URL https://www.
wandb. com/
Software available from wandb.com.
Dunean P Brown, Nandini Krishnamurthy, and Kimmen Sjölander. Automated protein subfamily
identification and classification. PLoS Comput Biol, 3(8)ee100, 2007.
Tom В Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler. Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are Few-Shot learners. May 2020.
SK Burley and GA Petsko. Weakly polar interactions in proteins. Advances in protein chemistry, 19:125-189, 1988.
Justas Dauparas, Haobo Wang, Avi Swartz, Peter Koo, Mor Nitzan, and Sergey Ovchinnikov. Uni- fied framework for modeling multivariate distributions in biological sequences. June 2019.
Jacob Devlin, Ming-Wor Weig Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. October 2018.
Yilun Du, Joshua Meier, Jerry Ma. Rob Fergus, and Alexander Rives. Energy-based Models
for Atomic-Resolution Protein Conformations. Technical report, sep 2019. URL https:
com httpab/ccco/accboooeeeeacco/o/
S D Dunn, M Wahl, and G B Gloor. Mutual information without the influence of phylogeny or en- tropy dramatically improves residue contact prediction. Bioinformatics, 24(3):333-344 Febrrurr 2008.
Richard Durbin, Sean R Eddy, Anders Krogh, and Graeme Mitchison. Biological sequence analysis: probabilistic models of proteins and nucleic acids. Cambridge university press, 1998.
Magnus Ekeberg, Cecilia Lövkvist, Yueheng Lan, Martin Weigt, and Erik Aurell. Improved contact prediction in proteins: using pseudolikelihoods to infer potts models. Phys. Rev. E Stat. Nonlin. Soft Matter Phys., 87(1)0012707, January 2013.
Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rihawi, Yu Wang, Llion Jones, Tom Gibbs, Tamas Feher, Christoph Angerer, Debsindhu Bhowmik, and Burkhard Rost. Prot- Trans: Towards cracking the language of life's code through sape-Supeee deep learning and high performance computing. July 2020.
WA Falcon. Pytorch lightning. GütHub. Note: http://siiiiaacooom///yyycciiitt lightning, 3, 2019.
Sarosh Fatakia, Stefano Costanzi, and Carson C. Chow. Computing Highly Correlated Positions Using Mutual Information and Graph Theory for G Protein-Coupler Receptors. PLoS ONE, 4 (3)e44611, mar 2009. ISSN 1932-6203.
doi: 10.137
0004681. URL https:
7/dx. www.ppp....o. org. 10. 1371/ journal. pone. C004681
Anthony A. Fodor and Richard W. Aldrich. On Evolutionary Conservation of Thermodynanic
Coupling in Proteins. Journal of Biological Chemistry, 279((1))190466-99995 apr 2004. ISSN 00219258. doi: 10.1074/jc.M442260200 15C23994/]
nlm. nih gov/
Ulrike Göbel, Chris Sander, Reinhard Schneider, and Alfonso Valencia. Correlated mutations and residue contacts in proteins. Proteins: Structure, Function, and Bioinformatics, 18(4):309-17, 1994. ISSN 10970134, doi: 10.1002/pro0.4401844022 URL attps: /pwpbmdddd.. ncbi. .nm.
ni niw.gov/82a 8208723/

---

Under review as a conference paper at ICLR 2021
John Ingraham, Vikas K Garg,
Regina Barzilay, and Tommi Jaakkola. Gen- design. In Neural Information Pro-
erative models for graph-based protein cessing 9711
Systems. MIT. 2019.
URL https://papera nips cc/ pc/pappr////
Rainer Jaenicke. Stability and stabilization of globular proteins solution. Journal of Biotechnol- ogy, 79(3):193-203, 2000.
L. Steven Johnson, Sean R. Eddy, and Elon Portugaly. Hidden Markov model speed heuristic and it- erative HMM search procedure. BMC Bioinformatics, 11(1):431, aug 2010. ISSN 14712105, doi: 211-7----11--1-4311 URL 2twp:///..ccce///000011111111111
wwt.eiicccc.cimmeeee.. biomedeent
David Jones and Shaun M Kandathil. High precision in protein contact prediction using fully convolutional neural networks and minimal sequence features. Bioinformatics, 34(19):3308- 2018.
Hetunandan Kamisetty, Sergey Ovchinnikov, and David Baker. Assessing the utility of coevolution- based residue-residue contact predictions in sequence- and structure-rich era. Proc. Natl. Acad. Sci. U. S. A., 11033)):5677-15577, September 2013.
T Korber, M Farber, D H Wolpert, and A S Lapedes. Covariation of mutations in the V3 loop of human immunodeficiency virus type envelope protein: an information theoretic analysis. Pruc. Natl. Acad. Sci. U. A., 9011))7776677,88 August 1993.
Alan Lapedes, Bertrand G. Giraud, LonChang Liu, and Gary D Stormo. Correlated mutations in models of protein sequences: Phylogenetic and structural effects. Lect. Notes Monogr. Ser., 33: 236-256, 1999.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. arXiv preprint arXiv: 03265, 2019.
Steve W. Lockless and Rama Ranganathan. Evolutionarily conserved pathways of energetic connectivity in protein families. Science, 28665438) 299, oct 1999. ISSN 00368075. doi: 10.11266/ciencee28666 www.:/eeenee22264348888 /Ripattps:/scieneee http:://wceeeeeeecccceceeemeeooo
abstract.
Amy X. Lu, Haoran Zhang, Marzyeh Ghassemi, and Alan Moses. Self-supervised contrastive learn- ing of protein representations by mutual information maximization. bioRxiv, 2020. doi: 10. 0001/2020.09.04.288999
11/10/2020.0 09. 04. 283929
A Madani, B McCann, N Naik, N S Keskar, and others. ProGen: Language modeling for protein generation. arXiv preprint arxiv, 2020.
Duccio Malinverni and Alessandro Barducci. Coevolutionary analysis of protein subfamilies by sequence reweighting. Entropy. (1(11):1127, 2019.
Bohdan Monastyrskyy. Daniel D' Andrea, Krzysztof Fidelis, Anna Tramontano, and Andriy Kryshtafovych. New encouraging developments contact prediction: Assessment of the casp 1l results. Proteins: Structure, Function, and Bioinformatics, 84:131-144, 2016.
Ananthan Nambiar, Simon Liu, Mark Hopkins, Maeve Heflin, Sergei Maslov, and Anna Ritz. Trans- forming the language of life: Transformer neural networks for protein prediction tasks. June 2020.
Sergey Ovchinnikov, Hetunandan Kamisetty, and David Baker. Robust and accurate prediction of rssidue-residue interactions across protein interfaces using evolutionary information. Elife, 3: e02030, May 2014.
10

---

Under review as a conference paper at ICLR 2021
Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Peter Chen, John Canny, Pieter Abbeel, and Yun Song. Evaluating protein transfer learning with TAPE. In Wallach, H Larochelle, A Beygelzimer. F dAlché-Bue E Fox. and R Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 9689-9701. Curran Associates, Inc., 2019.
Michael Remmert, Andreas Biegert, Andreas Hauser, and Johannes Söding. HHblits: Lightning-fast iterative protein sequence searching by HMM-HMM alignment. Nature Methods, 92):173-175, feb 2012. ISSN 15487091. doi: 10. 1038/mhh/mthh. 1818. URL htt.:tt:/pppeeiccctooneetttteee exg/caap9/groua http:/tosppppaoaaaaaaaaaaa anaalaiisss..............aaaiiia www.:ggggiggeeeeeeeeeee
tbm=
Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Demi Guo, Myle Ott, C Lawrence Zitnick, Jerry Ma, and Rob Fergus. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. August 2020.
Joerg Schaarschmidt, Bohdan Monastyrskyy, Andriy Kryshtafovych, and Alexandre MJJ Bonvin. Assessment of contact predictions in casp12: co-evolution and deep learning coming of age. Proteins: Structure, Function, and Bioinformatics, 86:51-66, 2018.
Stefan Seemayer, Markus Gruher, and Johannes Söding. CCMpred - Fast and precise prediction of protein residue-residue contacts from correlated mutations. Bioinformatics, 30(21):312-3313) may 2014. ISSN 14602059. doi: 10. 10.109//iooiinommatis///uss URL htttp::/pubbme::
nebi. n1m. nih. gov
Andrew W. Senior, Richard Evans, John Jumper, James Kirkpatrick, Laurent Sifre, Tim Green, Chongli Qin, Augustin Žídek, Alexander W.R. Nelson, Alex Bridgland, Hugo Penedones, Stig Petersen, Karen Simonyan, Steve Crossan, Pushmeet Kohli, David T. Jones, David Silver, Ko- ray Kavukcuoglu, and Demis Hassabis. Improved protein structure prediction using potentials from deep learning. Nature, 577(7792);70--710, jan 2020. ISSN 14764687. doi: 10. 10.1038/ s41586-019- 1923-7. URL /ttt:s/:/6aa..oo//00000///166661
doi
Amir Shanehsazzadeh, David Belanger, and David Dohan. Is transfer learning necessary for protein landscape prediction?, 2020.
Rojan Shrestha, Eduardo Fajardo, Nelson Gil, Krzysztof Fidelis, Andriy Kryshtafovych, Bohdan Monastyrskyy, and Andras Fiser. Assessing the accuracy of contact predictions in casp13, Pro- teins: Structure, Function, and Bioinformatics, 1018)): 2019.
Martin Steinegger, Markus Meier, Milot Mirdita, Harald Vöhringer, Stephan J. Haunsberger, and Johannes Söding. HH-suite3 for fast remote homology detection and deep protein an- notation. BMC Bioinformatics, si $12859-019-3019-7. URL 20(1):473, sep 2019. ISSN 14712105. doi: 10.1186/ ics bi omedcentral.....mmmmmo.ooo...... com/
Baris E Suzek, Hongzhan Huang, Peter McGarvey, Raja Mazumder, and Cathy Wu. UniRef: com- prehensive and non-redundant UniProt reference clusters. Bioinformatics, 23(10): 1282-1288,
May 2007.
Baris E Suzek, Yuqi Wang, Hongzhan Huang, Peter B McGarvey, Cathy H Wu, and UniProt Consor- tium. Uniref clusters: a comprehensive and scalable alternative for improving sequence similarity searches. Bioinformatics, 31(6):926 0221669922222 2015.
William R Taylor and Michael Sadowski. Structural constraints on the covariance matrix derived from multiple aligned protein sequences. PLoS One, 6(12):e28265, 2011.
John Thomas, Naren Ramakrishnan, and Chris Bailey Kelllog- Graphical models of residue cou- pling in protein families. IEEE/ACM Transactions on Computational Biology and Bioinformatics, 5(2):183-197, 2008.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E Ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I Guyon, U V Luxburg, Bengio. H Wallach, R Fergus. S Vishwanathan, and R Gamett (eds.), Advances in Neural Infor- mation Processing Systems 30, pp. 5998-6008 Curran Associates, Inc., 2017.
11

---

Under review as a conference paper at ICLR 2021
Evolutionary Tree
Multiple Sequence Alignment
Figure 8: The tree on the right depicts evolution of a protein family. The protein at the root the ancestral protein, and the five proteins at the leaves are its present-day descendants. The alignment on the left the corresponding Multiple Sequence Alignment of observed sequences.
Multiple Sequence
Padded Sequences
Alignment
Figure 9: MSA for sequences from Figure compared to padded batch of the same sequences.
J Vig, Madani, LR Varshney, Xiong, and others. Bertology meets biology: Interpreting attention
in protein language models. arXiv preprint arXiv, 2020.
Sheng Wang, Siqi Sun, Zhen Li, Renyu Zhang, and Jinbo Xu. Accurate de novo prediction of protein contact map by ultra-deep learning model. PLOS Computational Biology, 13(1):1-34, 2017. doi: 10. foii 100.5222244.. URL pabi 1005324
M Weigt, R A White, Szurmant, and others. Identification of direct residue contacts protein- protein interaction by message passing. Proceedings of the, 2009.
Jianyi Yang, Ivan Anishchenko, Hahnbeom Park, Zhenling Peng, Sergey Ovchinnikov, and David Baker. Improved protein structure prediction using predicted inter-residue orientations. bioRxiv, pp. 846279, 2019. doi: 10.1101/846279. URL https: httph::/wwwe: /wwp:////wwwwww/ www. http://wwwwwo....oo//ooeeeeeeeee
11.1100/84627991
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank J Reddi, Sanjiv Kumar, and Suvrit Sra. Why adam beats sgd for attention models. arXiv preprint
1911v:111.033194, 2019.
A APPENDIX
A.1 MULTIPLE SEQUENCE ALIGNMENTS
We demonstrate key concepts of protein evolution in Figure 8 On the left is phylogenetic tree. The leaves represent five observed proteins in a family, while the root represents their most recent common ancestor. The ancestral protein every observed sequence has one yellow and one blue amino acid on its ends, which lock together pinch off the loop. The amino acids in the middle of the loop exhibit considerable differences
loop and evolution preserved this loop structure. Thus
12

---

Under review as a conference paper at ICLR 2021
within the family, presumably leading to variations in function. This variation within the protein family is captured by the MSA on the right through its placement of gap characters (white squares). Compared to standard padding, shown in Figure 9 the placement of gap characters in the MSA ensures all blue and yellow amino acids lie in a single column. This signal is obfuscated by standard padding.
The problem of multiplying aligning a set of sequences has been extensively studied in Bioinfor- matics and Computational Biology (Durbin al. 1998). Given set of N sequences of average be generated by performing dynamic programming
length T., full multiple alignment array of size O(EN). This is computationally intractable; the common workaround is to perform iterative pairwise alignment to fixed reference sequence. The MSAs we use were generated with the HHSuite (Steinegger et 2019).
an
A.2 RECOVERING FACTORED ATTENTION FROM STANDARD ATTENTION
We show more precisely that factored attention can be recovered from standard attention one of two ways
1. Computing queries and keys from one-hot positional encodings and values from one-hot sequence embeddings.
2. Using learned positional embedding and sequence embedding with identity query, key, and value matrices.
The symmetrization operator symm is not applied to Transformer multihead attention, so we present slight variation of factored attention without symm in this section.
Single attention layer. Given a sequence of dense vectors x = -x11..., IL) with 2; E R", the attention mechanism of the Transformer encoder (multihead scaled dot product self-attention) pro- duces a continuous representation E IR Lxp, , If head size is d, this representation is computed using H heads (WQ, o, WK. Wy), where Wo, WK, Wy Rexd Queries, keys, and values defined as Q = zWQ, K = xWK. K. V = zWv. For single head (Wo. WK. Wv), the output is given by
QKT
y = softmax
V.
The full output in RdH is produced by concatenating all head outputs. A single Transformer encoder layer passes the output through dense layer, applying layer-norms and residual connection to aid optimization.
For the first layer, the input æ a sequence discrete tokens. To produce dense vector combining sequence and position information, positional encodings and sequence embeddings are combined. The positional encoding Eppoe E LRe produces a dense vector of dimension for each position dense vector
The sequence embedding Fzeg E RAxe maps each element of the vocabulary of dimension Typically these are combined through summation to produce a dense vector ă, (se((x;) + Epos (i), which is input to the Transformer as described above.
For this paper, we use only multihead self-attention without the dense layer, layernorm, or residual connections, as these drastically hurt performance when employed for one layer.
Factored attention from standard attention. Written explicitly, the input Transformer layer com- putes queries for a single head with Q = (Epos + Esen (z)) Wq. Keys and values are computed similarly. To recover factored attention, we instead compute queries and keys via Q Epo. Wo and K = WKx K. while values are given by V Exeg(z) W.. . For simplicity, we one-hot encode both position and sequence, which corresponds using identity matrices Epos = In E R f. and Eseg = IA E A Equivalently, one can view the positional encoding and sequence embedding as being learned while fixing Wo. WK, and Wy to be identity matrices.
Implicit single-site term in single-layer attention. For a single layer of attention, the product Epos Wy is matrix in L This matrix does not depend sequence inputs, thus allowing it to act as a single-site term. This suggests why inclusion of an explicit single-site term in F2gure22]jn had no effect for single-layer attention.
13

---

Under review as a conference paper at ICLR 2021
A.3 COMPUTATIONAL EFFICIENCY OF SINGLE LAYER MODELS
Model
Length
Batches/s
120
36.59
Potts
440
5.21
904
1.04
120
32.89
Factored attention (H = 256)
440
3,91
904
1.03
120
17.15
Single-layer attention
440*
1.685
(H 128)
904
0.46*
Table 1: Throughput of various models for batches of size 32. Stars (*) indicate usage of gradient accumulation due to GPU memory constraints. For sequences of length 440, batches size 16 were used for standard attention. For sequences of length 904, batches of size 4 were used.
In Table show the number of gradient steps per second for Potts, factored attention, and single- layer attention. We fix batch size of 32 and report numbers for families of lengths 120, 440, and 904 in order to explore the impact of length on computational efficiency for all models. Metries collected on node with a single NVIDIA RTX 2080Ti, an Intel i7-7800x with 6 physical cores and 12 virtual cores, and 32 GB of system RAM. To get steady state batches per second, model for 5000 steps and report the mean batches per second between steps 500 and 4500. Potts and factored attention have similar throughput, while standard attention scales much worse in length due to its memory requirements necessitating the use of gradient accumulation. This poor memory performance because standard attention must compute LX L attention maps for each batch element separately, whereas the I. x I. component of other models does not depend on sequence.
run each
A.4 LOSSES
The loss for all three models is of the form
((0:2) [((;2) cR(()),
where L is either pseudolikelihood masked language modeling and a regularizer.
Potts regularization. Consider the order-4 interaction tensor W, where W E RAX gives the parameters associated edge (i.j). We regularize W' by setting R(0)) Eiss || This term multiplied by A- L A, following Ovchinnikov al. (2014).
Factored attention regularization. Since factored attention is also a fully connected pairwise MRF, identical regularization to that of Potts. The order-4 tensor W is given by
: = Σ symm softmax (WAW (b((ab).
(4)
1=1
Single-layer attention regularization. Due the lack of MRF interpretation for single-layer at- tention, we chose to use weight decay as is typically done for attention models. This corresponds to setting R(#) be the sum of Frobenius norm squared for all weights WQ. Wx and Wy.
Single-site term. When any model has a single-site term, we follow standard practice and regularize its Frobenius norm squared.
14

---

Under review as a conference paper at ICLR 2021
A.5
PROTBERT BFD HEAD SELECTION
layer
head
P@ L
29
7
0.517
29
8
0.396
29
4
0.394
29
2
0.353
29
11
0.333
29
0
0.299
28
3
0.275
29
15
0.177
29
6
0.167
29
12
0.158
28
4
0.141
29
9
0.139
28
6
0.125
28
5
0.125
3
4
0.115
28
11
0.106
Table 2: The top 16 heads in ProtBERT-BFD whose attention maps gave the most precise contact maps across 500 validation families. Most of the top performing heads are found in the last layer. The top six heads were selected for our contact extraction in all results.
A.6 DATA AND METRICS
A.6.1 SELECTION OF PROTEIN FAMILIES
We used the following sets of families for model development:
1. A set of 748 families was chosen for performance evaluation. All metrics reported in the paper are on this set, with a single choice of hyperparameters for Potts models, factored attention, and standard attention. The 748 families were chosen randomly from the Yang jet al.) (2019) dataset, which consists of 15,051 MSAs generated from the databases Uni- Clust30 and UniRef100 (Suzek et al.) (007). as well metagenomic datasets. Our random sample is representative of the range of MSA depths and protein lengths, see Figure 10
2. A set of six families from the 748 was chosen to choose hyperparameters for single-layer attention. They chosen span range of MSA depth (large and small), well as three different regimes of Potts performance (Good, Ok, Poor). These families were used tune hyperparameters as described in Section A.7.1 See Table
3. A set of ten families from the 748 chosen where factored attention performed very poorly in our initial experiments. Half were chosen to be long proteins and the other half be short. This set used optimize learning rate and regularization for factored attention to ensure reasonable model performance. See Table 4
4. 500 entirely disjoint families were further selected randomly from Yang et al.) (2019) and used to compute average precision for each head in ProdBERT-BFD (Elnaggar al. 2020 Performance on these families was used for selecting the top heads, see Table
PDB
Sequences
Length
Potts Performance
3er7 I A
33673
118
Good
5fo5 B
17560
88
Ok
2w 2a_18_14.
33619
308
Poor
4gnr_
2073
351
Good
5mkc. A.A
515
207
Ok
1_9
146
292
Poor
Table 3: 6 families chosen for hyperparameter optimization for single-layer attention.
15

---

Under review as a conference paper at ICLR 2021
PDB
Sequences
Length
_
2145
140
1_A
5535
143
3cy4_1_B
1064
154
6
2325
155
3p6b B
4353
186
1jml
17130
202
4yt2_1_A
15481
343
3vmm__A
4383
471
4ege_1_A
9929
539
3gq7-1_A 6568
605
Table 4: families chosen for hyperparameter optimization for factored attention
Training data distribution
1000
trRosetta dataset
training data
800
Length of reference sequence
600
400
200
0
27
28
29
210
211
22
213
24
# of sequences in MSA
Figure 10: The length and MSA size distribution for our 748 family subset (red) compared the full 15,051 families in the trRosetta dataset selected for training
A.6.2 PRODUCING CONTACT MAPS
A PDB structure gives 3D coordinates for every atom in structure. We Euclidean distance between the beta carbons to define distance between any pair positions. A pair of positions where this distance is less than 8. is declared to be contact.
16

---

Under review as a conference paper at ICLR 2021
A.6.3 SCORING CONTACT PREDICTIONS
Given a predicted contact map Õ E _gL and a true contact map C' {0, 10,11 we describe metrics for scoring C.
sequence x (#1,,..,,71, vr.) length L has (5) potential contacts. Since we see (L) contacts, contact prediction a sparse prediction task. Accordingly, we focus on precision-recall based quanti- tative analyses of C. Common practice in the field is sort all (2) entries of Ò in decreasing order and evaluate precision at various length thresholds, such as the top L or L/10 predictions (Shresthal et al. 2019). Note that this analysis is similar to choosing recall cutoffs along precision-recul curve, where sorted length index plays the role of recall on the axis. Unlike recall, length-based based cutoffs do not rely on knowledge of the actual number of contacts. In addition to the precision at various length (recall) cutoffs, also computed Area Under the Precision Recalll Curve (AUC). AUC is a widely used metric for comparing classifiers when the positive class rare.
A.7 HYPERPARAMETERS
Potts. We used À = 0.5, learning rate of 0.5, and batch size 4096. Pad, gap, and mask were all encoded with the same token.
The Potts model is trained using a modified version of Adam presented in Dauparas et al. 2019 This modification
made to improve performance of Adam to match that of L-BFGS.
Factored attention. We AdamW with learning rate of 5e-3 and set X = 0.01. The default head size was set 32 unless stated otherwise.
Single-layer attention. We set embedding size of 256, head size 64, and number of heads 128. The model is trained with Adam W using learning rate of 5e-3 and weight decay of 2e-3. Attention dropout of 0.05 is also applied. The batch size is 32 and mask prob for masked language modeling 0.15. We use separate mask token and pad.gap token.
ProtBERT-BED. ProBEERT-BFD has 30 layers each with 16 heads and hidden size of 1024. The training dataset is mixture of UniRef50 (Suzek al. 2015) and BFD, It has 2, 122 million protein sequences. See Elnaggar et al. 2020) more information.
A.7.1 HYPERPARAMETER SWEEP DETAILS
Potts. The Potts model implementation using psuedolikelihood has been optimized by others, not tune performance. Since performance with MLM was comparable to pseudolikelihood, we did not sweep for MLM either.
Single-layer attention. Standard attention is by far the most sensitive model to hyperparameters. To find a reasonable set of hyperparameters, we first swept the six families in Table 3 performing a grid search over
H E (32, 64, 128, 256, 512)
d (32, 64, 128, 256, 512)
e (128, 256, 512}
attention dropout in (0, 0.05,.05000 0.1}
learning rate in (5e 3, 2}
3, -3} weight decay in (0,
We found that the choice Н 256, d 64, 256, attention dropout 0.05, learning rate of 5e - 3 and weight decay of 2e 3 performed well across all six families. Due to GPU memory constraints, we had set H 128 for further runs.
Factored attention. We swept factored attention over the families in Table performing grid search over
learning rate in {1e 3, 5e 3, 2,5e 2}
17

---

Under review as a conference paper at ICLR 2021
regularization coefficient in {le 4, be 4, 3, 3, le 1e - 2]
We found that learning rate of 5e-3 and regularization of 0.01 were effective, but that other configu-
rations such as regularization of 5e-3 also performed well. Both H and are evaluated extensively in our results.
A.8 ADDITIONAL FIGURES
6000
5000
Total Number of Contacts
4000
3000
2000
1000
250
500
750
1000
1250
1500
1750
2000
Protein Length
Figure 11: The total number of contacts for a structure as a function of protein length follows linear trend. (slope 2.64, R² 0.929)
18

---

Under review as a conference paper at ICLR 2021
1.0
Proportion of Residues in trRosetta Dataset
0.8
0.6
0.4
0.2
0.0
0
1
2
3
4
13 Number of per-Residue Contacts
5
6
7
9
10
11
12
14
15 16 17 18
Figure 12: The empirical CDF of number of per-residue contacts for 3,747,101 residues in 15,051 structures in the trRosetta dataset.
1.0
0.8
Precision a
0.6
0.4
0.2
0.0
8
16
32
64
128
256
512
Potts
Number of heads
Figure 13: Reducing the number of heads causes much steeper decrease in precision L.
19

---

Under review as a conference paper at ICLR 2021
1.0
Spcarman
Pearson
0.4
0.0
Number
Figure 14: Effect of number of heads on correlation between the order-4 weight tensors for factored attention (see Equation 4) and Potts (see Section 3)
0.5
0.-
0.2
03
250
Abrrie ReadS
Albuicn HealSar
Figure 15: Effect head size on factored attention precision at and L/5 over 748 families. Increasing head size has small effect precision, though not nearly pronounced the number of heads.
P@L L 0.48
P@L L 0.73
100
200
300
400
H
Potts
Figure heads has degraded performance for precision at L.
20

---

Under review as a conference paper at ICLR 2021
Factored Attn Frozen Values
1.00
16384
# of MSA sequences
0.75
4096
0.50
1024
0.25
256
0.00
0.0
0.5
1.0
Potts
Figure 17: Factored attention trained with single set of frozen value matrices performs comparably to Potts, evaluated on precision at across 748 families.
156
Potts
H=256
Parameters (Million)
100
H=64
H=8
56
25
6
0
0
200
400
600
800
Length of reference
Figure 18: Number of parameters versus length for MRF models.
21

---

Under review as a conference paper at ICLR 2021
1.0
APC
False
0.8
True
0.6
Pr@0 1
0.4
0.2
0.0
Potts
Factored Alin
Atn
PoxBERT-BFD
Figure 19: APC has a significant positive effect on the performance of Potts and factored attention. It makes only slight difference on the performance of the other two models.
10384
10384
Fastord Ann MLM
0.8
of MSA soqjAaTiONS
0.8
SOQUAATIONNS
4096
Posts M
-4006
0.4
1024
0.4
1024
0.0
0.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
Pscodolikelibood
Pscudolikclihood
Figure 20: Effect of loss precision at L over many families. Pseudolikclihood has uniform but small benefit over masked language modeling for both models.
1.0
Aligned Lagut
Fabse
Thue
0.5
0.5
ALC
0.
0.2
Pactord Aucation
Soanerd Aucation
Figure 21: Training on unaligned families degrades performance on almost all families.
22

---

Under review as a conference paper at ICLR 2021
1.0
Fulse
085
True
0.6
ALC
0.-
0.2
0.0
Factored Aucation
Standird Aucction
Figure 22: The addition of a single-site term to either factored or standard attention produces little additional benefit.
23


Reviews:
[{"review_id": "uBqGHjM1FQb", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This manuscript describes a connection between Potts models and attention as implemented in modern transformers. The authors then present an attention model in which positional encodings are defined as one-hot vectors indicating fixed positions in the multiple sequence alignment and train single layer attention models. These models, unsurprisingly, perform similarly to Potts models without APC correction for contact prediction. The methods section is somewhat confusingly written. I think the factored attention model would benefit from being described on it’s own terms rather than in connection with typical multiheaded attention, especially because the isolation of position encodings and amino acids at those positions dramatically simplifies the understanding of W_Q, W_K,  and W_V. The authors also spend a long time describing well known methods, but without providing additional insight. The connection between the Potts model and attention described in this paper should be obvious to those who already understand attention models and Potts models and the empirical results of the factored attention model don’t make this approach seem compelling. In the discussion, the authors make several broad future speculations. Some of these would be interesting contributions and I encourage the authors to develop this work further. Maybe factored attention could be promising for better capturing dependencies between positions for deeper transformers on MSAs, but it isn’t likely that this work will be of broad interest to the machine learning community. This manuscript seems better suited to a workshop or other specialized venue. Some specific comments on this work follow below.\n1.\tIn the factored attention model, the authors use one-hot encoding of the position index as the position encoding. This is equivalent to learned position embeddings as in BERT which is worth mentioning. \n2.\tThe authors discuss single-site potentials as a difference between Potts models and single layer attention models and then show a comparison of attention models with and without single-site potentials showing little difference. However, attention models already implicitly have single-site potentials which arise from the positional encoding input features. Granted, this is not the case for the factored attention model where single-site potentials seem to have more effect, though in the negative direction.\n3.\tThe authors state that “The ability of factored attention to capture similar contacts to Potts without use of APC suggest that it may be more suitable for protein design.” I don’t follow this conclusion. If the factored attention model performs equivalently to the Potts model alone and worse than the Potts model with APC correction, why would it be more suitable for protein design?\n4.\t What makes the single-layer attention or factored attention models compelling for protein modeling? What problems do these models solve that are not better solved by the Potts model or traditional transformers?\n\nWhat would raise my score:\n1.\tPresent a compelling use case for the factored attention model. What questions can be answered (or better answered) with this model over the Potts model or other alternatives? One idea is to use the factored attention model as the layers in a full deep transformer model and see if this architecture can improve tasks where MSA training data is available.\n\nEdit: I have increased my score in light of the response and manuscript edits. The manuscript is improved, but I think the method still needs more development. There are a number of interesting pieces but the final picture of an improved protein model is not fully resolved.", "extra_scores": {}, "extra_text": {}}, {"review_id": "AowgLbK53ne", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary\n=======\nTransformers models have been recently shown to capture protein contact information in their attention maps when trained unsupervised of millions of protein sequences. This paper draws parallels between Transformers and Potts models (fully-connected pairwise MRF)--the current standard approach for protein contact prediction--and shows empirically that Transformers are competitive with Potts models. Understanding the differences and similarities between Transformers and Potts models makes Transformers less of a ‘black-box’ and helps to establish them as a principled method for contact prediction. The paper is clearly written and the evaluation is solid. I have only a few comments.\n\n\nMajor comments\n=============\n1. What is the maximum sequence similarity between the training sequence of ProtBERT and sequences in TrRosetts alignments that were used for testing? Sequences must not overlap have a maximum similarity of let’s say 80%.\n\n2. You describe that you used three sets of families from the TrRosetta dataset (A.4.1). Why did you use only 732 families for testing (set 3)? Were these all families that were not included in the first two sets? How many families do the first two sets include and how similar are families of different sets? Ideally, train, tune, and test families belong to different super families.\n\n3. You describe in section A.3 how you extracted protein contact maps from the attention maps of ProtBERT. This is an important detail that must be described in the main text. How did you choose the 6 heads? Did you choose them manually or, for example, by training a linear model to predict contacts from attention maps and using the weights for identifying important heads, or computing the weighted average of attention maps?\n\n\nMinor comments\n=============\n4. Section 3.2, ‘x = E_seq(x_i) + E_pos(i)’: How did you compute positional embeddings and why do and add embeddings instead of concatenating them?\n\n5. Section 3.2, ‘We treat the positional embedding E_pos as an overall summary per-position information’. Please describe more clearly what this summary is.\n\n6. Section 4, first paragraph: The L of the precision at L metric is not the sequence length but the number of top sequences. You describe L as being both. \n\n7. Figure 6 is not discussed. Instead of showing this figure, I suggest quantifying the correlation depending on the number of heads by computing and discussing  the Spearman correlation.\n\n8. Rives et al  2020 ‘Biological structure and function emerge…’ have recently shown in addition to Vig et al that protein contact can be predicted from attention maps, which must be also pointed out in the ‘Background’ section.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "A7p-BK8Hhd", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Recently, some researchers tried to apply attention models into the protein field, using self-supervised learning to predict protein contacts. In this work, the author attempt to build the connection between such works and the old-school model, Potts model. By simplifying some operations within the attention model, the author managed to build an analog between the simplified model and the Potts model. The analog is intuitive and easy-to-understand. The authors further compare the simplified model and the Potts model on 748 protein families, showing that they are similar. Or probably the simplified attention model is even better. This is an interesting work. However, I also have a number of concerns. The advantages and disadvantages are listed below.\n\nPros:\n1. The manuscript is concise and easy-to-understand.\n2. The idea is intuitive and reasonable, with experimental support.\n\nCons:\n1. The analog between the simplified attention model and the Potts model is intuitive but not rigorous. The authors claim that they provide a theoretical connection between the two models. However, that part is not strong enough, without proof.\n2. There are two assumptions in this work, which make the simplified model different from the attention models that the previous researchers used. Firstly, they train the model on multiple sequence alignment instead of the raw sequences. If they train the model on the raw sequences, the performance is unacceptable, as shown in Figure 16, which is consistent with the previous research. Secondly, they removed the sequence embedding in queries and keys. This simplification makes the model only consider the statistical pattern in the MSA. To me, this one is a too strong assumption. \n3. The running time and hardware comparison is missing. If the single layer of attention is comparable to the Potts model, not outperform it significantly, while it would take much more time to train, the researchers would need to think twice if they want to use the attention model. \n4. The ablation study makes me feel that the results are on the opposite of the conclusion. Here is my logic. With the above two assumptions, the attention model can achieve similar performance as the Potts model, or a little bit better. However, when we train on the unaligned sequences, which is the usual case that we would use the attention model, the performance becomes unacceptable. Then why we want to use the more expensive attention model? The attention model in the NLP field is a different story. Those models are refreshing the STOA performance all the time. However, in the protein field, the attention model can still only achieve comparable performance as the classific models, after a two-year study. They seldom outperform classic algorithms. The results in this manuscript are consistent with the previous research. So I am not convinced regarding the conclusion in the abstract:\n\"Taken together, these results provide motivation for training Transformers on large protein datasets.\"\n5. The potential audience of this paper would be those who are specialized or interested in bioinformatics and protein.", "extra_scores": {}, "extra_text": {}}, {"review_id": "35c7F7Aganm", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThis paper explores the connection between the classic Potts model-based approaches and modern Transformer-based approaches for protein contact map prediction. To this end, the authors introduce a simplified variation of the attention layer called factored attention, and show that a single layer of factored attention performs operations similar to those performed by the Potts model-based methods.\n\nPros:\n- The paper attempts to connect classic and modern approaches to protein contact map prediction, which might be interesting to the people working in this field. The evidence presented (simplifying attention layer so that the equations look similar to the classic methods, numerical results of the simplified attention layer close to the classic methods) is reasonably convincing.\n- The topic of the paper is quite timely, there has been a lot of interest recently in modelling proteins using the latest NLP techniques.\n- The paper is well written. I appreciate the effort put in by the authors to define basic protein terminologies which might not be obvious to readers without biology background. \n\nCons:\n- The contributions of the paper would have been more interesting if the proposed modifications of the attention layer led to increased prediction performance of models which are representative of the state-of-the-art. Specifically, if retraining ProtBERT-BFD using the modified attention layer led to further improvement in performance, that would have been a solid contribution.\n- Are MRF models really that competitive for contact map prediction? From what I understand, deep neural networks have been far better at this task for quite some time now. At multiple places in the paper, the authors give the impression that MRF models are close to state-of-the-art.\n- In the last paragraph of the introductory section, the idea of encoding the MSAs is introduced which seemed interesting. However, from what I understood from the rest of the paper, the queries and keys are extracted solely based on the position of the amino acid. Is that right? If so, does the position correspond to the position in the sequence or in the MSA? Are the actual alignments used in any of the results in the paper? Please clarify.\n\nComments:\n- Section 3.1: \"each edge\" should have a capital e.\n- Section 3.3, specifically the part where you show that factored attention is a pairwise MRF, is too brief. Given that this is a main contribution of the paper, it would be worthwhile to explain this connection in a more detailed manner.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper shows a connection between Potts model and Transformers and uses the connection to propose a factored attention energy to use in an MRF. Results are shown, using this energy based on factored attention. Also, pretrained BERT models are used to predict contact maps as a comparison.
The reviewers found the paper interesting from a protein structures prediction point of view, but from a machine learning perspective their opinion was that the paper does not offer a coherent, compelling method that is very novel, and the connection between Potts and an energy based attention model is not that overwhelming.  In addition the presentation was somewhat circuitous.  

The authors made improvements to the paper over the course of the review, which is appreciated, but the method presented does not match the target for an ICLR paper in terms of methodological contributions.  
Decision: Reject