Neural Potts Model
Authors: Tom Sercu
Robert Verkuil
Joshua Meier
Brandon Amos
Zeming Lin
Caroline Chen
Jason Liu
Yann LeCun
Alexander Rives
Venue:  (2021)
Keywords: proteins
potts model
unsupervised learning
amortized optimization
structure prediction

Abstract: We propose the Neural Potts Model objective as an amortized optimization problem. The objective enables training a single model with shared parameters to explicitly model energy landscapes across multiple protein families. Given a protein sequence as input, the model is trained to predict a pairwise coupling matrix for a Potts model energy function describing the local evolutionary landscape of the sequence. Couplings can be predicted for novel sequences. A controlled ablation experiment assessing unsupervised contact prediction on sets of related protein families finds a gain from amortization for low-depth multiple sequence alignments; the result is then confirmed on a database with broad coverage of protein sequences.

Under review as a conference paper at ICLR 2021
NEURAL POTTS MODEL
Anonymous authors Paper under double-blind review
ABSTRACT
We propose the Neural Potts Model objective as an amortized optimization problem. The objective enables training single model with shared parameters explicitly model energy landscapes across multiple protein families. Given a protein sequence input, the model is trained predict a pairwise coupling matrix for Potts model energy function describing the local evolutionary landscape of the sequence. Couplings can be predicted for novel sequences. A controlled ablation experiment assessing unsupervised contact prediction on sets of related protein families finds a gain from amortization for low-depth multiple sequence alignments; the result is then confirmed on a database with broad coverage of protein sequences.
INTRODUCTION
When two positions in a protein sequence are in spatial contact in the folded three- three-dimensiona structure of the protein, evolution is not free choose the amino acid at each position independently. This means that the positions co-evolve: when the amino acid at one position varies, the assignment at the contacting site may vary with it. A multiple sequence alignment (MSA) summarizes evolutionary variation by collecting a group of diverse but evolutionarily related sequences. Patterns of variation, including co-evolution, can be observed in the MSA. These patterns are in turn associated with the structure and function of the protein (Göbel et al., 1994). Unsupervised contact prediction aims to detect co-evolutionary patterns in the statistics of the MSA and infer structure from them.
The standard method for unsupervised contact prediction fits Potts model energy function to the MSA (Lapedes al., 1999; Thomas et al., 2008; Weigt et al., 2009). Various approximations are used in practice including mean field (Morcos et al., 201 1), sparse inverse covariance estimation (Jones et al., 2011 1), and pseudolikelihood maximization (Balakrishnan et al., 2011; 1; Ekeberg et al., 2013; Kamisetty et al., 2013). To construct the MSA for a given input sequence, similarity query is performed across a large database to identify related sequences, which are then aligned to each other. Fitting the Potts model to the set of sequences identifies statistical couplings between different sites in the protein, which can be used to infer contacts in the structure (Weigt et al., 2009). Contact prediction performance depends the depth of the MSA and is reduced when few related sequences are available to fit the model.
In this work we consider fitting many models across many families simultaneously with parameter sharing across all the families. We introduce this formally as the Neural Potts Model (NPM) objective. The objective an amortized optimization problem across sequence families. A Transformer model is trained to predict the parameters of a Potts model energy function defined by the MSA of each input sequence. This approach builds on the ideas in the emerging field of protein language models (Alley et al., 2019; Rives et al., 2019; Heinzinger et al., 2019), which proposes to fit single model with unsupervised learning across many evolutionarily diverse protein sequences. We extend this core idea to train model to output explicit energy landscape for every sequence.
To evaluate the approach, we focus on the problem setting of unsupervised contact prediction for proteins with low-depth MSAs. Unsupervised structure learning with Potts models performs poorly when few related sequences are available (Jones et al., 20 Kamisetty et al., 2013; Moult al., 2016). Since larger protein families are likely to have structures available, the proteins of greatest interest for unsupervised structure prediction are likely to have lower depth MSAs (Tetchner et al., 2014). This is especially a problem for higher organisms. where there are fewer related genomes (Tetchner et al., 2014). The hope is that for low-depth MSAs, the parameter sharing in the neural model will improve results relative to fitting independent Potts model each family.

---

Under review as a conference paper at ICLR 2021
32
Search & align
MSA
Optimize Potts model
W
(b)
x
NPM forward pass
(a(
Figure 1: (a) Standard Potts model requires constructing an MSA and optimizing parameters W. (b) Neural Potts Model (NPM) predicts W in a single feedforward pass from single sequence.
We investigate the NPM objective in a controlled ablation experiment on a group of related protein families in PFAM (Finn et al., 2016). In this artificial setting, information can be generalized by the pre-trained shared parameters to improve unsupervised contact prediction on subset of the MSAs that have been artificially truncated to reduce their number of sequences. We then study the model in the setting of a large dataset without artificial reduction, training the model on MSAs for UniRef50 sequences. In this setting there is also an improvement on average for low depth MSAs both for sequences in the training set as well as for sequences not in the training set.
BACKGROUND
Multiple sequence alignments MSA is set of aligned protein sequences that are evolutionarily related. MSAs constructed by retrieving related sequences from sequence database and aligning the returned sequences using a heuristic. An MSA can viewed as a matrix where each row is a sequence, and columns contain aligned positions after removing insertions and replacing deletions with gap characters.
Potts model The generalized Potts model defines a Gibbs distribution over a protein sequence (Z1,...,TL)) of length L with the negative energy function:
E(x) E (z,)) Je(())))))
Which defines potentials h2 for each position in the sequence, and couplings Jij for every pair of positions. The parameters of the model are W = {h, J} the set of fields and couplings respectively. The distribution p(æ; ) is obtained by normalization
E(w; W)} (WW.
Since the normalization constant is intractable, pseudolikelihood is commonly used to fit the parame- ters (Balakrishnan et al., 2011; I; Ekeberg al., 2013). Pseudolikelihood approximates the likelihood of sequence as product of conditional distributions: (p W) = = -e log p(z; x= W). To estimate the Potts model, take the expectation:
(Pl(W)=
= E (x:
2~
MSA M. In practice we have finite set of sequences M in the MSA estimate Eq. (2). La regularization p(W)= A, ||ll|2 + An. ||A]]2 added, and sequences reweighted to account for redundancy (Morcos et al., 2011) 1). We write the regularized finite sample estimator as:
(P)(WW
Mof Σ
"T W)) + p(W)
Which sums over all the M sequences the finite MSA M, weighted with Th summing collectively Mat. The finite sample estimate of the parameters W obtained by minimizing Lri.
Idealized MSA Notice how in Eq. (2), we idealized the MSA M as a distribution, defined by the protein family. We consider the set of sequences actually retrieved in the MSA M in Eq. (3) as a finite sample from this underlying idealized distribution. For some protein families this sample will contain more information than for others, depending on what sequences are present in the database. We will refer to W as hypothetical idealized estimate of the parameters to explain how the Neural Potts Model can improve the finite sample estimate W for low-depth MSAs.
2

---

Under review as a conference paper at ICLR 2021
2.1 AMORTIZED OPTIMIZATION
We review amortized optimization (Shu, 2017), generalization of amortized variational inference (Kingma Welling, 2013; Rezende et al., 2014) that uses learning to predict the solution to continuous optimization problems to make the computation more tractable and potentially generalize across problem instances. We are interested in repeatedly solving expensive optimization problems
(x)) = arg min L(W; W
(4)
where W E Ih is the optimization variable, I E I is the input conditioning variable to the optimization problem, and L: R.m x IR + the objective. We assume (*xx) is unique. We consider the setting of having distribution over optimization problems with inputs x ~ p(z), and the arg min of those optimization problems (zz).
Amortization uses learning to leverage the shared structure present across the distribution, e.g. solution W"(x) is likely correlated with another solution W W' (x'). Assuming an underlying regularity of the data and loss L, we can imagine learning to predict the outcome of the optimization problem with an expressive model (x such that hopefully Wa s W Modeling and learning Wa(r) are the key design decisions when using amortization.
Modeling approaches. In this paper we consider models Wo(r) that directly predict the solution to Eq. (4) with neural network, approach which follows fully amortized variational inference models and the meta-learning method (Mishra et al., 2017). The model can also leverage the objective information <r) and gradient information Ww (r: e.g. by predicting multiple candidate solutions W and selecting the most optimal one. This is sometimes referred to as semi- cemi-mmortization or unrolled optimization- based models and is considered in Gregor & LeCun (2010) for sparse coding, Li & Malik (2016) 6); Andrychowicz et al. (2016); Finn et al. (2017) for meta-learning, and Marino et al. (2018); Kim al. (2018) for posterior optimization.
Learning approaches. There are two main classes of learning approaches for amortization:
arg min E L(Wo(a);)) p(z)
(5)
arg min E (W((=)-"(<<)))
(6)
9
# p(z)
Gradient- based approaches leverage gradient information of the objective and optimize Eq. (5) whereas regression-buseer approaches optimize distance to ground-truth solutions W*, such as the squared L distance in Eq. (6). Prior work has shown that models trained with these objectives can learn to predict the optimal W directly as function of Given enough regularity of the domain, if we observe new (test) samples x' ~ p(x)) we expect the model to generalize and predict the solution to the original optimization problem Eq. (4). Gradient- based approaches have the computational advantage of not requiring the expensive ground-truth solution W while regression- less susceptible to poor local optima in L. Gradient- based approaches are used in variational inference (Kingma & Welling, 2013), style transfer (Chen & Schmidt, 2016), meta learning (Finn al., 2017; Mishra et al., 2017), and reinforcement learning, e.g. for the policy update in model-free actor-critic methods (Sutton & Barto, 2018). Regression based approaches are more common in control for behavioral cloning and imitation learning (Duriez et al., 2017; Ratliff al., 2007; Bain Sammut, 1995).
approaches
NEURAL POTTS MODEL
In Eq. (2) we introduced the Potts model for a single MSA M (aligned set of sequences x), to
optimize W = (h" J arg miny Ea~M (4r. W)]. As per Eq. (5) We will now introduce neural network to estimate Potts model parameters from a single sequence: (ha( Ja( = Wz( with single forward pass.
We propose minimizing the following objective for the NPM parameters 0, which directly minimizes the Potts model losses in expectation over our data distribution æ D and their MSAs き M(z):
ENPM(()
E M((=)
(ā Wo(x))
a>D
3

---

Under review as a conference paper at ICLR 2021
(P) Training objective
Generalization
Amortization gap
L(W) loss
" "underfittin"
L(W*)
Inductive gain
[(W)
W
Wo
W
W {
Inductive generalization gain (illustration with a ID loss landscape). W is the standard Potts model, estimated on the finite observed MSA M. Though it minimizes the training objective, does not achieve perfect generalization performance. However the Neural Potts Model Wo(x) can generalize better than W through transfer learning from related samples. guided by the inductive bias of the model. We expect this especially when the estimate W far from W e.g. on small or biased MSAs.
Figure 2:
To compute the loss for a given sequence x we compute the Potts model parameters (x and evaluate its beeudo-ikelihood loss fp1 query sequence. This fits exactly in "amortized optimization" in Section 2.1 Eq. (5); we train a model predict the outcome of set of highly related optimization problems. One key extension the described amortized optimization setup is that the model Wo estimates the Potts Model parameters from only the MSA query sequence as input rather than the full MSA M(z). Thus, our model must learn to distill the protein energy landscape into its parameters, since it cannot look up related proteins during runtime. full algorithm given Appendix
of sequences E from the MSA constructed with as
Similar to the original Potts model, we need to add regularization penalty p(W) to the main objective. For a finite sample N different query sequences {æn}. and a corresponding sample of Nx M aligned sequences { from MSA M(æn), the finite sample regularized loss, i.e. NPM training objective, becomes:
M Σ
(N) (0)=
Σ
w (<r" Wo(m.))) + ((Wo((z.
Mett nn)
Inductive generalization gain (see Fig. 2) is when the Neural Potts Model improves over the individual Potts model. Intuitively this is possible because the individual Potts Models are not perfect estimates (finite/biased MSAs), while the shared parameters of We can transfer information between related protein families and from pre-training with another objective like masked language modeling (MLM).
Let us start with the normal amortized optimization setting, where we expect an amortization gap (Cremer et al., 2018). The amortization gap means that Wo(r) will be behind the optimal W for the objective L: (x)) > L(W*). This closely related to underfitting: the model W, is not flexible enough to capture (x:). However, recall that in the Potts model setting, there is finite-sample training objective È (Eq. (8)), with minimizer W We can expect amortization gap in the training objective; however this amortization gap can now be advantageous. Even the amortized solution Wo(x)) (2x is near-optimal on É, it can likely find a more generalizable region of the overparametrized domain W by parameter sharing of 6 allowing it to transfer information between related instances. The inductive bias of Wa(x) can allow the neural amortized estimate to generalize
better, especially when the finite sample M is poor. This inductive bias depends on the choice of model class for We. its pre-training, as well as the shared structure between the protein families in the dataset. Concretely loss on test MSA sequences, but also the performance on downstream validation objectives like predicting contacts, a proxy for the model's ability to capture the underlying structure of the protein.
will consider for the generalization loss L not just the eseudo-likelihow
4

---

Under review as a conference paper at ICLR 2021
We will show that for some samples (x) L(W*), i.e there is inductive generalization gain. This is visually represented in Fig. 2; and Table compares amortized optimization and NPM, making connection to multi-task learning (Caruana, 1998). Additionally, we could frame NPM as hypernetwork, a neural network that predicts the weights of second network (in this case the Potts model) in, e-g., Gomez & Schmidhuber (2005); Ha al. (2016); Bertinetto al. (2016).
Table 1: Comparison between (A) "standard" amortized optimization, (B) Neural Potts Model, and (C) Multi-task learning. From row (A) amortized optimization to (B) Neural Potts Model, a finite-sample training loss is introduced which comes with considerations of generalization and regularization. This is related to multi-tast learning, but with a major difference that (B) the solo optimization is over single tensor W in the Potts model, but (C) function So in learning problem. In the amortized/multi- role that different related tasks play in (C) MTL. In the NPM setting (B), Wg takes explicitly as argument, versus (C) MTL typically just has a separate output head per task.
setting, the distribution over query sequences a: in (B) NPM plays the
Solo objective
Solo objective Generalization
Amortized / Multi-task
Parametrization model choices
Training
(A)
Solo: W E R
C(s; W)
C(s; W)
Amortized optim:
Optim-
Amor: Wa IR R"
(= Training)
Ep(s) C(s; Wo(s))
Amortized
+learner class
(B)
PLL. finite MSA
Distr L(W) B[ (Ā; W)]
Neural Potts
Solo: W E R"
Potts
M: L(W)= Σ (p W)
Ex L (We(x)))
+regularization
NPM
or
Amor: We R IR"
Contact pred
+learner class
(C)
ERM: L(fe) Σ ((fo(Im)) ym)
C(fe) [.((fo()) v) ay
Multi-task learning: (C= for T related
Solo: fo R +
ML MTL
+regularization
(Multi-task
+learner class
learning)
MTL: f3 R2d IR
tasks
+ param sharing fó
In summary, the goal for the NPM is to "distill" an ensemble of Potts models into single feedforward model. From self-supervised learning perspective, rather than supervising the model with the input directly, we use supervision from an energy landscape around the input.
EXPERIMENTS
In Section 4.1 we present results on small set of related protein domain families from Pfam. where we artificially reduce the MSA depth for few families study the inductive generalization gain from the shared parameters, In Section 4.2 we present results on large Transformer trained on MSAS for all of UniRef50.
For the main representation (xx) we use a bidirectional transformer model (Vaswani et al., 2017). 7). To compute the cour-dimensional pairwise coupling tensor Jxx)) from sequence embedding ge (x) we introduce the multi-head bilinear form (mhbf) in Appendix B. One can think of the multi-head bilinear form as the x self-attention maps of the Transformer's multi-head attention module, but without softmax normalization. When using mhbf for direct prediction, there are K heads, one for every amino acid pair k, I. For the Pfam experiments, we extend the architecture with convolutional layers after the mhbf, where the final convolutional layer has K2 output channels. We initialize go (a:) with Transformer pre-traine pre-trained with masked language modeling following (Rives et al., 2019).
To evaluate Neural Potts Model energy landscapes, will focus proteins with structure in the Protein Data Bank (PDB), using the magnitude of the couplings after APC correction to rank contacts. The protocol is described in Appendix C.2.
5

---

Under review as a conference paper at ICLR 2021
Range Short
Range Medium
Long
0.20
Toget: Proctaton
0.15
0.3
0.100
Mothed
Indep, Potts
0.075
0.10
NN Potts
0.030
0.05
0.1
0.025
0.000
0.00
161
102
1.0
10
1.0
10 10:
10a
10*
MSA Degth
MSA Depth
MSA Depth
Figure 3: Contact prediction precision on Pfam families from the NADP Rossmann clan, at different levels of depth reduction. Columns show (from left to right) short, medium and long-range precision for top-L. threshold. Across the metrics, NPM outperforms the independent Potts model trained on the shallowest MSAs, as well as the Nearest Neighbor Potts model baseline.
Gradient Updatos
Precision Top 1. Long Range
10010
0100
2000
0.25
3000
4000
5000
0.20
MSA Depth
0.15
1000
150
200
250
400
450
NPM amortized pscado--ikellhood
Figure 4: Trajectory of training on the NADP Rossman clan, averaged over five-fold cross-evaluation.
4.1 PFAM CLANS
To study generalization in controlled setting, we investigate a small set structurally-rellted MSAs from the Pfam domain family database (Finn al., 2016) belonging the same Pfam clan. We expect that on a collection of related MSAs, information could be generalized to improve performance on low-depth MSAs. Families within Pfam clan are linked by distant evolutionary relationship. giving them related but not trivially-similar structure. We obtain contact maps for the sequences in each of the families where structure is available in the PDB. At test time input the sequence and compare the generated couplings under the model to the corresponding structure.
We compare the NPM two baselines. The first direct comparison is to independent Potts model trained directly on the MSA. For the second baseline we construct the "nearest neighbor" Potts model, by aligning each sequence against all families in the training set, and using the Potts model from the closest matching family.
We perform the experiment using a five-fold cross-evaluation scheme, in which we partition the clan's 's families into five equally-sized buckets. As in standard cross-validation, each bucket will eventually serve as an evaluation set. However, we do not remove the evaluation bucket. Instead, we artificially reduce the number of sequences in the MSAs in the evaluation bucket to smaller fixed MSA depth. MSAs in the remaining buckets remain unaltered. The goal of this setup is to check the model's ability to infer contacts artificially limited sets of sequences. Both NPM and the baseline independent Potts model are fit on the reduced set of sequences. Note that while the baseline Potts model uses the reduced MSA of the target directly, NPM is trained the reduced MSA but evaluated using only the target sequence as input. We train a separate NPM on each of the five eross-evaluation rounds, evaluate the structures corresponding the bucket with reduced

---

Under review as a conference paper at ICLR 2021
MSAs, and show averages and standard deviations across rounds. Further details are provided for model training in Appendix C.1 and for the Pfam dataset in Appendix C.3.
Figure 3 shows the resulting contact prediction performance the 181 families in the NADP Rossmann clan, with additional results on the P-loop NTPase, HTH, and AB hydrolase clans in Appendix D Fig. 9. We initialize 12-layer Transformer with protein language modeling pre-training. Because of the small dataset size, we keep the weights of the base Transformer Og frozen and only finetune the final layers. As a function of increasing MSA depth, contact precision improves for both NPM and independent Potts models. For the shallowest MSAs, NPM has higher precision relative the independent Potts models. The advantage at low MSA depth is most pronounced for long range contacts, outperforming independent Potts models up to MSA depth 1000. These experiments suggest NPM able to realize an inductive gain by sharing parameters in the pre-trained base model as well as the fine-tuned final layers and output head. Figure 4 shows training trajectories. We observe near-monotonic decrease of the amortized in the evaluation set, and increase of the top-L long range contact precision. This indicates that improving the NPM objective improves the unsupervised contact precision across the reduced-depth MSAs. Furthermore worse contact precision.
leuudo-lkeeihood loss (Eq. (7)) on the MSAs
expected overfitting for smaller MSA depth: better training loss but
Additionally, assess performance of different architecture variants: direct prediction with the multi-head bilinear form (always using symmetry), with or without tied projections, and addition of convolutional layers after the multi-head bilinear form. The variants described in detail in Appendix B. We find in Appendix D Fig. 8 that addition of convolutional layers after the multi-head bilinear form performs best; for the variant without convolutional layers, the head without weight tying performs best.
4.2
UNIREF50
Kerge Modon
Ragge Leng
Top-U/SLM precisions
4000
3.4
Top 15 (0-AAMM
1000
0.4
Tralo
2000
43
1938
Nookl Inday Potts
18
1000
NPM
2.6
03
1500
2000
Text
0.5
1990
1008
9.1
02
2.0
0.15
65. 90)
(50, 188)
(1.0. 590)
:500
(0.15)
(5. 90)
(TR, 180)
(102, 5000
2516
0.3
0.
1.0
Mal
MaTf
Indepeateer Port
Figure 5: UniRef50: contact prediction precisions (higher is better) with 95% bootstrapped confi- dence intervals, on medium range (left), long range (middle), binned by MSA depth MaR. Top row: sequences from the train set; bottom row: sequences from the test set. For shallow MSAs, average performance of NPM higher than the independent Potts model. Right: seatter plot comparing long range precision from NPM vs independent Potts model, each point is a protein. More metrics are presented in Appendix Fig. 10.
We now perform an evaluation in the more realistic setting of the UniRef50 dataset (Suzek et al., 2007). First we examine MSA depth across UniRef50 (Suzek al., 2007). Appendix C.4 Fig. 7 finds that 19% of sequences in UniRef50 have MSAs with fewer than 10 sequences. (38% when a minimum query sequence coverage of 80% is specified).
7

---

Under review as a conference paper at ICLR 2021
We ask whether amortization gain be realized in two different settings: (i) for sequences the model has been trained on: (ii) for sequences in the test set. We partition the UniRef50 representative sequences into 90% train and 10% test sets, constructing an MSA for each of the sequences. During training, the model is given a sequence from the train set as input, and the NPM objective is minimized using sample from the MSA of the input sequence. In each training epoch, a different set of 30 sequences from the MSA to fit the NPM objective. We use ground truth structures evaluate the NPM couplings and independent Potts model couplings for contact precision. The dataset is further described in Appendix C.4; and details on the model and training are given in Appendix C.I.
randomly subsample
The independent Potts model baseline is trained on the full MSA. This means that in setting (i) the NPM and independent Potts models have access to the same underlying MSAs during training. In setting (ii) the independent Potts model is afforded access to the full MSA; however the NPM has not been trained on this MSA and must perform some level of generalization to estimate the couplings.
Figure 5 shows comparison between the NPM predictions and individual Potts models fit from the MSA. The Neural Potts Model is given only the query sequence as input. On top-L/5 long range precision, NPM has better precision than independent Potts models for 22.3% of train and 22.7% of test proteins. We visualize Fig. 6 example proteins with low MSA-depth where NPM does better than the individual Potts model. For shallow MSAs, the average performance of NPM is higher than the Potts model, suggesting an inductive generalization gain.
To contextualize the results let consider the setting where the amortized Neural Potts Model (i) matches the independent Potts model on training data: this means the NPM can predict good quality couplings from single feedforward pass without access the full MSA at inference time; (ii) surpasses the independent model on training data: the amortization helps the NPM to improve over independent Potts models, i.e. realizes inductive generalization gain; (iii) matches the independent model on test sequences: indicates the model is able to synthesize a good Potts model for sequences not in its training data; (iv) surpasses the independent model on test sequences: the model actually improves over an independent Potts model even for sequences not in its training data. In combination these results indicate a non-trivial generalization happens when NPM trained on UniRef50.
Train Sample, PDB ID: 3a0u.
Test Sample, PDB ID: 2ejx.
20
10
NPM True Positive
60-
False
Indep. Potts True Positive Indep. Potts Palse Positive
80-
100
100
320
NPM Top-L/5 LR:
0.87
NPM Top-L/5 LR: 0.56
Indop. Potts Top-L/5 LR:
Indep. Potts Top-L/5 0.11
M. eff:
68
M eff:
1087
Figure 6: Examples where NPM outperforms the independent Potts model fit directly on the MSA. NPM top-L/5 LR contact prediction (lower diagonal, red) compared to the independent Potts model prediction (upper diagonal, blue). All ground truth contacts are indicated in black. True and false hits are indicated with dots and crosses, respectively.
RELATED WORK
Recently, protein language modeling has emerged as promising direction for learning representations of protein sequences that are useful across a variety of tasks. Rives et al. (2019) and Rao et al. (2019) trained protein language models with the masked language modeling (MLM) objective originally proposed for natural languge processing by Devlin et al. (2019). Alley et al. (2019), Heinzinger et al. (2019), and Madani al. (2020) trained models with autoregressive objectives. Transformer protein
8

---

Under review as a conference paper at ICLR 2021
language models trained with the MLM objective learn information about the underlying structure and function of proteins including long range contacts (Rives et al., 2019; Vig et al., 2020). This paper builds the ideas in the protein language modeling literature, introducing the following new ideas: the first is supervision with an energy landscape (defined by a set of sequences) rather than objectives which are defined by single sequence; the second is use amortized optimization fit a single model across many different energy landscapes with parameter sharing; the final is the consideration of the unsupervised contact prediction problem setting rather than the use of representations in supervised pipeline.
Unsupervised structure learning reviewed in the introduction. The main approach has been to learn a set of constraints from family of related sequences by fitting a Potts model energy function to the sequences. Our work builds on this idea, but rather than fitting a Potts model to a single family of related sequences, proposes through amortized optimization to fit Potts models across many sequence families with parameter sharing in a deep neural network.
Supervised learning has produced breakthrough results for protein structure prediction (Xu, 2018; Senior et al., 2019; Yang et al., 2019). State-of-the-art methods use supervised learning with deep residual networks on co-evolutionary features derived from the unsupervised structure learning pipeline. While Xu et al. (2020) show that reasonable predictions can be made without co-evolutionary features, their work also shows that these features contribute significantly to the performance of state-of-the-art pipelines.
Prior work studying protein language models for contact prediction focuses the supervised setting. Bepler & Berger (2019) studied pre-training an LSTM on protein sequences and fine-tuning on contact data. Rives et al. (2019) and Rao al. (2019) studied supervised contact prediction from Transformer protein language models. Vig et al. (2020) found that contacts are represented in Transformer self-attention maps. Our work differs from prior work on structure prediction using protein language models by focusing on the unsupervised structure learning setting. It would be logical extension of this work integrate the Neural Potts model into the supervised pipeline.
DISCUSSION
This paper explores how protein sequence model can be trained to produce local energy landscape that is defined by set of evolutionarily related sequences for each input. The training objective is cast amortized optimization problem. By learning to output the parameters for Potts model energy function across many sequences, the model may learn to generalize across the sequences.
We also formally and empirically investigate the generalization capability of models trained through amortized optimization. We consider the setting of training independent Potts models on the MSA of each sequence, in comparison with training single model using the amortized objective predict Potts model parameters for many inputs. Empirically the amortized objective provides an inductive gain when few related sequences are available in the MSA for training the independent Potts model.
A number of direct extensions exist for future work, including further investigation of model ar- chitecture and parameterization of the energy function by the deep network, use of the amortized models in supervised pipeline, and combining independent Potts models with amortized couplings. The hidden representations could also be investigated for structure prediction and other tasks using the approaches in the protein language modeling literature. The main contribution of this work to directly incorporate information from set of sequences related to the input in the learning objective. It would be interesting to investigate other possible approaches for incorporating this type of supervision into models that aim to learn underlying structure from sequence data.
REFERENCES
Ethan C Alley, Grigory Khimulya. Surojit Biswas, Mohammed AlQuraishi, and George M Church. Unified rational protein engineering with soquence-only deep representation learning. bioRxiv, pp. 589333, 2019.
Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, and Nando de Freitas. Learning leam by gradient descent by gradient descent. arXiv preprint 166600444 2016.

---

Under review as a conference paper at ICLR 2021
Michael Bain and Claude Sammut. A framework for behavioural cloning. In Machine Intelligence 15, pp. 103-129, 1995.
Sivaraman Balakrishnan, Hetunandan Kamisetty, Jaime G Carbonell, Su-In Lee, and Christopher James Lang- mead. Learning generative models for protein fold families. Proteins: Structure, Function, and Bioinformatics, 79(4):061-1107, 2011.
Tristan Bepler and Bonnie Berger. Learning protein sequence embeddings using information from structure. arXiv:/902.0666) 2019.
Luca Bertinetto, João F Henriques. Jack Valmadre, Philip Torr, and Andrea Vedaldi. Learning feed-forward one-shot learners. In Advances in neural information processing systems, pp. 523 531. 2016.
Rich Caruana. Multitask learning. In Learning learn, pp. 95-133. Springer, 1998.
Tian Qi Chen and Mark Schmidt. Fast patch-based style transfer of arbitrary style. arXiv preprint vrxiv:661..04333 2016.
Chris Cremer, Xuechen Li, and David Duvenaud. Inference suboptimality in variational autoencoders.
arXiv:1801.03533 2018.
Jacob Devlin, Ming-Wo Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In Proceedings of the 2019 Conference of the North (A/merican Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume / (Long and Short Papers), pp. 4171-4188 Minncapolis, Minnesota, 2019. Association for Computational Linguis- tics. doi: 10.1 10118863///NN 19- 1423. URL wwwthttpp///ii..otreee/ abs 1810. 04805.
Stanley D Dunn, Lindi M Wahl, and Gregory Gloor. Mutual information without the influence of phylogeny or entropy dramatically improves residue contact prediction. Bioinformatics, 24(3)333-340, 2008
Thomas Duriee, Steven L Brunton, and Bernd Noack. Machine learning control-taming nonlinear dynamics and turbulence, volume 116. Springer, 2017.
Magnus Ekeberg, Cecilia Lövkvist, Yueheng Lan, Martin Weigt, and Erik Aurell. Improved contact prediction in proteins: using pseudolikelihoods to infer potts models. Physical Review E, 87(1):017777 2013.
Chelsca Finn, Pieter Abbecl, and Sergey Levine. Model-agnostic meta-lcarning for fast adaptation of deep
networks. arXiv preprint arXiv: 034000000, 2017.
Robert D Finn, Penelope Coggill, Ruth Y Eberhardt, Sean R Eddy, Jaina Mistry, Alex L Mitchell, Simon C Potter, Marco Punta, Matloob Qureshi, Amaia Sangrador Vegas, et al. The pfam protein families database: towards a more sustainable future. Nucleic acids research, 44(D1):D279 -D255, 2016.
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-31) 317, 1994.
Faustino Gomez and Jürgen Schmidhuher. Evolving modular fast- networks for control. International Conference on Artificial Neural Networks, pp. 383-389. Springer, 2005.
Karol Gregor and Yann LeCun. Learning fast approximations of sparse coding. In Proceedings of the 27th International Conference on Machine Learning (ICML-10), pp. 399-406, 2010.
David Ha, Andrew Dai, and Quee V Le. Hypernetworks. arXiv:1100.00)) 09100000990990...... 2016.
Michael Heinzinger, Ahmed Elnaggar, Yu Wang, Christian Dallago, Dmitrii Nechaev, Florian Matthes, and Burkhard Rost. Modeling aspects of the language of life through transfer-leaming protein sequences. BMC bioinformatics, 20(1):723, 2019.
David Jones, Daniel WA Buchan, Domenico Cozzetto, and Massimiliano Pontil. Psicov: precise tural contact prediction using sparse inverse covariance estimation on large multiple sequence alignments. Bioinformatics, 184-199884999 2011.
Hetunandan Kamisetty, Sergey Ovehinnikov, and David Baker. Assessing the utility of coevolution residue- residue contact predictions in a sequence-and structure-rich era. Proceedings of the National Academy of Sciences, 110739)
2013.
Yoon Kim, Sam Wiseman, Andrew C Miller, David Sontag, and Alexander M Rush. Semi-amortized variational 2018.
autoencoders.
10

---

Under review as a conference paper at ICLR 2021
Diederik P Kingma and Max Welling, Auto-encoding variational bayes. ICLR, 2013.
Alan S Lapedes, Bertrand G Giraud, LonChang Liu, and Gary D Stormo. Correlated mutations in models of
protein sequences: phylogenetic and structural effects. Lecture Noter-Monograph Series, pp. 236-256, 1999.
Ke Li and Jitendra Malik. Learning optimize. urXiv:1606.01885. 2016.
Ali Madani, Bryan McCann, Nikhil Naik, Nitish Shirish Keskar, Namrata Anand, Raphael R. Eguchi, Po-Ssu Huang, and Richard Sacher. ProGen: Language Modeling for Protein Generation. arXiv, 3 2020. URL http:/ / aaxiiv wttpi//.rgga//b//////e 200220000003497..
Joseph Marino, Yisong Yue, and Stephan Mandt. Iterative amortized inference. arXiv: 2018.
Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel. A simple neural attentive meta-leamer. arXiv:1707.03144 2017.
Faruck Morcos, Andrea Pagnani, Bryan Lunt, Arianna Bertolino, Debora S Marks, Chris Sander, Riceardo Zecchina, José N Onuchic, Terence Hwa, and Martin Weigt. Direct-coupling analysis of residue cocvolution captures native contacts across many protein families. Proceedings of the National Academy of Sciences, 108 (49):E1293-E130)) 2011.
John Moult, Krzysztof Fidelis. Andriy Kryshtafovych, Torsten Schwede, and Anna Tramontano. Critical assessment of methods of protein structure prediction: Progress and new directions in round xi. Proteins: Structure, Function, and Bioinformatics, 84:4-14, 2016.
Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Chen, John Canny, Pieter Abheel, and Yun Song. Evaluating protein transfer learning with tape. arXiv:1906.00 08230, 2019.
Nathan Ratliff, J Andrew Bagnell, and Siddhartha Srinivasa. Imitation learning for locomotion and ma- nipulation. 2007 7th IEEE-RAS RAS International Conference on Humanoid Robots, pp. 392-397. IEEE, 2007.
Danilo Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In Proceedings of the 31st International Conference on Machine Learning (ICML-14), pp. 1278-1228 1286, 2014.
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 250 million protein sequences. bioRxiv, pp. 622803, 2019. URL /tttps://doi.ooo 1101/ 622803.
Stefan Seemayer, Markus Gruber, and Johannes Söding. Compred fast and precise prediction of protein residue-residue contacts from correlated mutations. Bioinformatics, 30(21):3 11:1288-31300 2014.
Andrew W. Senior, Richard Evans, John Jumper, James Kirkpatrick, Laurent Sifre, Tim Green, Chongli Qin, Augustin Zidek, Alexander W. R. Nelson, Alex Bridgland, Hugo Penedones, Stig Petersen, Karen Simonyan, Steve Crossan, Pushmeet Kohli, David T. Jones. David Silver. Koray Kavukcuoglu, and Demis Hassabis. Protein structure prediction using multiple deep neural networks in the 13th Critical Assessment of Protein Structure Prediction (CASP13). Proteins: Structure, Function, and Bioinformatics, 87(12):111--1148 12 2019. ISSN 0887-3585. doi: 10. 10.1002/prot.258.. URL httptttppp...o / in ibbrry wi Tey com/ doi 8oi/abs/10 10000/000ooo/ 25834.
Rui Shu. Amortized optimization. http://ttshiii/00///////nnaaaaat 2017. Retrieved 2020-
09-22.
Martin Steinegger and Johannes Süding. Clustering huge protein sequence sets in linear time. Nature communi- cations, 9( (1):1-8, 2018
Martin Steinegger, Markus Meier, Milot Mirdita, Harald Vöhringer, Stephan J Haunsberger, Johannes Söding. Hh-suite3 for fast remote homology detection and deep protein annotation. BMC bioinformatics, 20 (1):1-15, 2019.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Baris E. Suzek, Hongzhan Huang. Peter McGarvey, Raja Mazumder, and Cathy H. Wu. UniRef: Comprehensive 2007. ISSN 13674803,
and non- redundant UniProt reference clusters. Bioinformatics, 23(10):1282-13 doi: 10. 1933/oonffomati////m000. URL http: /twp://wwww.//.////////////tt. un prot org. UniRef50 database licensed under (CC BY 4.0).
11

---

Under review as a conference paper at ICLR 2021
Stuart Tetchner, Tomasz Kosciolek, and David T Jones. Opportunities and limitations in applying coevolution- derived contacts to protein structure prediction. Bio-Algorithms and Med- hedd-Sytems, 10(4):243-255 2014.
John Thomas, Naren Ramakrishnan, and Chris Bailey-Kellogg Graphical models of residue coupling in protein families. IEEE/ACM Transactions on Computational Biology and Bioinformatics, 5(2):183- 2008
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention you need, NeurIPS, 2017.
Jesse Vig, Ali Madani, Lav R Varshney, Caiming Xiong. Richard Socher, and Nazneen Fatema Rajani. Bertology meets biology: Interpreting attention in protein language models. arXiv:2006. 15222, 2020.
Martin Weigt, Robert White, Hendrik Szurmant, James A Hoch, and Terence Hwa. Identification of direct residue contacts in protein protein interaction by message passing. Proceedings of National Academy of Sciences, 106(1):67-72, 2009.
Jinbo Distance-based protein folding powered by deep learning. arXiv:1811003481, 2018.
Jinbo Xu, Matthew Mcpartion, and Jin Li. Improved protein structure prediction by deep leamnngg irrespective of co-evolution information. bioRxiv, 2020.
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.
12

---

Under review as a conference paper at ICLR 2021
Appendix
A LEARNING THE NEURAL POTTS MODEL
Algorithm Learning the Neural Potts Model
Initialize the amortized model We while unconverged do
Sample N sequences &n 'D and M MSA sequences T.t ~ M(xx) grad_upde _upd_deeae(, VoLNrM (0; (xn), (# end while
& Optimize the regularized loss in Eq. (8)
B MODEL ARCHITECTURE: MULTI-HEAD BILINEAR FORM FOR PAIRWISE COUPLINGS
In this Section, we describe the model architecture to compute four-dimensional pairwise coupling tensor Jn (x) from sequence embedding go (z))
B. MULTI-HEAD BILINEAR FORM
We write sequence length L and amino acid vocabulary K 21. The single site potentials BLXK and the pairwise couplings are a four-dimensiona tensor: J E 4x x indexed as Jiy(k, 1).
We start with sequence-level model produce the embedding e of the sequence (typically final hidden layer output): e = (x) E IR Ixxd The estimator for single-site potential ha(x) is linear projection layer on the embedding; (x) go (x) with ph Rd Now we discuss how to parametrize the estimator Jo(æ)) E =
Multi- head bilinear form for direct prediction We introduce multi-head bilinear form (mhbf) on the embedding e; i.e. for every pair k. of amino acids we have a bilinear form, parametrized with learned interaction matrix Bed E dddd connecting the hidden states at positions c,2 ej E IR xxd So we compute the K2 bilinear forms for amino acid pairs (k,1) between L x I. position pairs (i,j): (k,(k, e,Bble We always use a low-rank decomposition Bkl = kkl y kiT with both
Tikl yul E TRdxd', so the bilinear form becomes the inner product in the lower-dimensional space with d' the projection dimension: (e;U*) V We can interpret this as an inner product of embeddings i after linear projection to space specific amino acid pair (I. This low-rank multi-head bilinear form is similar to the multi-head attention mechanism introduced in Vaswani et al. (2017), but without softmax normalization.
Notation-wise, our parameters 0 include the parameters of the transformer that produces the embed- ding and the components of the interaction matrix kU ykg.
Direct prediction: tied projection One way reduce the number of parameters in the multi muat bilinear form, is for the low-rank decompositions of the K2 heads Bbl to share their decomposition per k. I. Instead we can share/tie the projection matrices between amino acids k, I: U U and ku Vl. such that head B U Note that the dot product in this case is after a linear projection specific to single-site amino acid k and separately; Jej = e, = (e,U*) (e,"")..
Direct prediction: Symmetry We can or should parametrize the estimator of J to be symmetrical against interchanging both i.j and k.l: Jiy (k, (L. k)), i.c. no difference between the order of considering interaction between AA lo at position with AA I at position j. This does not mean symmetry each interaction matrix! We ask that
(k, BB =e, B. BltT
13

---

Under review as a conference paper at ICLR 2021
The second equality is the symmetry, the last equality by transposing the bilinear form. From Bhd = B follows that
Ukipki Villglh T - -il
The last equality is the obvious choice. In the tied parametrization, this simply becomes U U vk ylk such that Wa = U Once again, note that the dot product now becomes (k, = (e,U") We present a Tensor decomposition perspective on this multi-head bilinear form in Appendix B.2.
Convolutional layers after multi-head bilinear form As extended model architecture, consider having convolutional layers after the multi-head bilinear form (only used for the Pfam experiments), parametrized with learned interaction matrix B E Rdxd In this case, rather than having K2 heads Bki we now have an arbitrary number of heads F which will become the number of channels the consecutive convolutional layers: B U VFF. We add 1 convolutional layers having also F channels, and finally K2 output channels for the last convolutional layer. Weight tying and symmetry considerations of the mhbf do not apply this model variation.
B.2 TENSOR DECOMPOSITION VIEW ON MULTI-HEEA BILINEAR FORM
We can the multi- bilinear form tensor decomposition of J, for which will use Einstein notation to indicate that any pair of indices appearing both in subscript and superscript are summed over their range. Let us write the tensor collecting the U matrices as U E K xdxd', and index into in the same notation for U: Uut U'kl With a, [1 [l...dl d], r E [1 [1...dd The same for V. Now the J estimate in the full untied asymmetric case, written as tensor, becomes
J c U Valr of
or the symmetric (4 Vitk) and tied (4 E Kxdxd version:
J = e U ar ug ef
Note that the 24 V are shared across proteins, while the embeddings = ge(or) are specific per protein, based on a high-capacity sequence level model.
C EXPERIMENT DETAILS
C.1 TRAINING DETAILS
We summarize the precise model architecture and optimization settings in Table 2. During each NPM training step, for a given input æ, M sequences Tt are randomly sampled (M=100 or 30, see Table 2), for the pseudo-likelihood loss evaluation in Eq. (8). Each sequence is selected with probability according to its sequence weight wm. One can think of these sampled sequences as similar minibatch. Note that compute the independent Potts model baseline, the Potts model is computed without any downsampling of the MSA. Additionally, in the Pfam experiments the loss term for family n in Eq. (8) is upweighted with a factor / (Mr(()) which places more weight on the well-formed, deep MSAs and discounts the shallower MSAs. In both the Pfam and UniRef experiments, we enforce max sequence length of 512 via random contiguous crops of positions.
C.2 VALIDATION DETAILS
To compute precisions, we convert the pairwise couplings J E [ to an L x L pairwise coupling score by (1) zeroing all positions in J corresponding to gap characters, (2) computing the magnitude via Frobenius norm over the K x K matrix Jag for every pair of positions i, and (3) applying the Average Product Correction (Dunn al., 2008). True contacts are defined pairwise distances less than or equal to Angstroms. Precision calculated as the true positive fraction of the top L, [./2], [L/5] predicted contacts. Additional to precision, the Area Under the
14

---

Under review as a conference paper at ICLR 2021
Pfam
UniRe
Parameters
111M (frozen:
519M
26M trained)
Number of layers
12
32
Embed dim
768
I024
Model details
Attention heads
12
16
mhbf U, projection di-
128
512
mension
mhbf number of heads
128
K = 441
Conv after mhbf mhbf tied
x 1 (3 layers)
N/A
No
of epochs
21
50
MLM pre-training (UniRef50)
of gradient updates
175k
449k
Learning rate
4e-4
3e-5
Number updates
10k
135k
Adam learning rate
le-4
3e-5
p(W) multipliers A, An
le-3, le-2
le-3, le-2
NPM training
M (samples per MSA)
100
30
Batch size (# seqs)
64
2048
Max sequence length
512
512
Base model frozen
Yes
No
Table 2: Hyperparameters
Precision-Recal Curve (AUC) computed, summing over thresholds stepwise per L/10 increment up to L. Precision and AUC metrics are computed at sequence separations of short (6 < < 12), medium (12 S < 24). and long (24 s) ranges.
For the independent Potts model baselines in all experiments, we use CCMpred (Seemayer et al., 2014), GPU implementation of pseudolikelihood maximization (Balakrishnan et al., 2011). The coupling matrix J from the independent Potts model is processed in the same way following steps (1-3) described above.
C.3 PFAM TRAINING DATA AND SETUP
Data Selection. We use the Pfam database (Finn et al., 2016) version 28.0. All MSAs in the HTH (n=217), P-loop NTPase (n=198), NADP Rossmann (n=181), and AB hydrolase (n=67) clans were parsed from the multiple alignment file Pfam-A.full. We apply two preprocessing steps to all MSAs. First, for speed, we only load up to maximum of 100k sequences from each MSA. Next, we apply HHfilter, from the HHSuite3 (Steinegger et al., 2019) toolset, with all default settings, to each MSA. We find that filtering improves contact prediction accuracy of the independent Potts model baseline.
Dataset splits. We perform the experiment using five-fold cross-evaluation scheme, in which we partition the clan's families into five equally gizze buckets. As standard cross-validation, each bucket will eventually serve as an evaluation set. However, we do not remove the evaluation bucket, but artificially reduce the number of sequences in the MSAs in the evaluation bucket to a small fixed MSA depth (=purging the MSA). All Pfam experiments are repeated 5 times, each with different selection for the reduced bucket. In our figures, we plot average results, with confidence interval bounds defined by the standard deviation across the five-fold cross-evaluation.
During NPM training, we iterate over the the set of MSAs the four buckets that have not been reduced, well as the reduced bucket. At given training step, randomly sclect sequence within an MSA for use as input to NPM. This selection is likely to return a sequence with inserted gap characters. We drop these gap characters and their corresponding columns in the MSA. Then we randomly subsample 100 sequences z from the MSA to fit the NPM objective. The procedure is described in more detail in Appendix C.I.
Evaluation During evaluation, we assess the NPM and the independent Potts model via a contact prediction task (described in previous subsection), the families in the evaluation bucket. For
15

---

Under review as a conference paper at ICLR 2021
CDF of MSA Depth
100
hhfilter cov=80%
hhfilter cov=60%
hhfilter cov=40%
80
hhfilter cov=20%
No hhfilter applied
Percentlle
60
40
20
10"
10
102
10a
Depth
Figure 7: A random sample of 1000 MSAs for 1000 sequences UniRef50 were analyzed. The graph shows cumulative density plots with MSA depth in log-scale the axis, for different query sequence coverage requirements. (-cov [20, 80]] ) specified by a call to HHfilter, applied to each MSA. The fraction of MSAs with depth I 10 is 19% (38% when coverage of 80% is specified), while the fraction of MSAs with depth < 100 is 30% (55% when a coverage of 80% is specified).
each family, a single structure is selected as target, using the pdbmap included in Pfam. NPM's contact predictions are made using only the sequence belonging to the target structure. To compute the independent Potts model for given family in the evaluation bucket, the depth-reduced MSA is aligned to the sequence from the target structure, and the Potts model is computed without further downsampling.
As an additional baseline, we predict contacts for validation sequences using the Potts model of the "Nearest Neighbor" family in the train set. For a given validation sequence, we calculate "nearness" to all train families via calls to HHalign given the sequence and the train family's 's Pfam seed alignment as input. We select the family with the highest HHalign probability score as the nearest neighbor. The nearest neighbor prediction generated as follows: (1) the validation sequence is aligned to the selected train family's MSA; (2) independent Potts model is fit to the selected train family's MSA, yielding predicted contact map for the train family; (3) the rows and columns of the predicted contact map that align to the validation sequence extracted to construct prediction for the validation sequence.
C.4 UNIREF50 TRAINING DATA AND SETUP
For the experiments in Section 4.2, we retrieve the UniRef-50 (Suzek al., 2007) database dated 2018- 03. The UniRef50 clusters partitioned randomly in 90% train and 10% test sets. For all sequences, we construct MSAs using HHblits (Steinegger et al., 2019) against the UniClust30_200 1113000221_1111 database. HHblits is using the default settings, for iterations with an e-value of 0.001.
is important to note that given this MSA generation procedure, validation sequences can be included in MSAs of train sequences. However, as input to NPM.
guaranteed that validation sequences are not trained on
Evaluation of contact precision
We use contact precision as a proxy to measure unsupervised structure learning in the underlying Potts model. To define a set of structures for evaluation, collect structures from the PDB, and assign them to either the training sequences or test sequences. This allows us to separately examine performance of NPM sequences from its training set, and sequences from its test set. Note, that
16

---

Under review as a conference paper at ICLR 2021
the structures used only to evaluate unsupervised contact prediction performance of the model; the model is never trained on structures.
We query the Protein Data Bank (PDB) to obtain list of all protein structures with resolution less than 2.5 Å, length greater than 40 residues, and a submission date before May 1, 2020. We search each pdb entry for hits against the sequences in the training and test sets for NPM respectively. the PDB entry retrieves hits only training sequences
assign it to the wraining-squeett the PDB entry retrieves hits only to test sequences we assign to the test-sequences group. Any PDB entry which hits both training and test sequences neither, is discarded. To perform the search use the MMSeqs2 software suite (Steinegger & Söding, 2018) using the default settings with 50% sequence identity at 80% target coverage. We then cluster each of the two groups of PDB entries 50% sequence similarity, resulting in a dataset of 1040 structures assigned to train sequences and 211 structures assigned validation sequences. MSA construction for the PDB entries precisely follows the procedure for UniRef50 (first paragraph); the method for contact prediction from the model couplings (for NPM or the independent Potts model) is described in Appendix C.2.
group. If
D ADDITIONAL EXPERIMENTS
17

---

Under review as a conference paper at ICLR 2021
0.30
Precision
0.15
NPM model
mhbl (tied) mhidd (untied)
0.30
mbilf ooew
0.05
Short,
Short, 1000
Medicmn, (Range, MSA Depth)
Medium, 1000
Long,
Long, 1000
Figure 8: Comparison of the main NPM model architecture choices, evaluated using the Pfam experimental setup. We show precisions fixed top-L threshold, while on the x-axis we vary sequence separation range and two levels of MSA depth reduction (10 and 1000). Standard deviations over the five-fold eross-evaluation are shown. For the direct multi-head bilinear form (mhbf) prediction (tied or untied), we found an improvement from using U, projection dimension 512 rather than 128. Other hyper-parameters follow Table 2.
Range
Short
Rasgee Modium
Long
0.12
10pt procitoon
0.10
0.25
P-loop NiPase
0.08
0.08
0.20
0.06
0.06
0.04
0.04
0.10
0.02
0.05
0.150
0.20
Top- Predicion
0.15
aapee Rossmann
NADD
0.100
0.075
0.0050
0.05
0.025
Method
0.000
0.00
0.0
Indep, Potts
0.20
NPM
0.20
0.150
NN Potts
Top-1 Precidoo
0.125
0.15
0.15
0.10
0.30
0.075
0.05
0.050
0.05
0.00
0.000
0.00
0.20
0.4
0.125
Teg-1 Preciston
0.100
0.15
Saw.aooeeoo 1or ITV
0.075
6.050
0.05
0.025
164
102
1.2
104
10
102
10a
10
10
102
10a
10*
MSA Depth
MSA Depth
MSA Depth
Figure 9: Contact prediction precision Pfam families at different levels of depth reduction. Three approaches are compared: (1) the Neural Potts Model, (2) independent Potts Model, (3) Nearest Neighbor Potts. On shallow MSAs, NPM outperforms both independent Potts and Nearest Neighbor Potts for all four of the clans. On deep MSAs NPM matches the independent Potts baseline for 3 of the 4 clans, Clans are noted on the right hand side.
18

---

Under review as a conference paper at ICLR 2021
0.5
04
044
0.2
Twp. prccions
0.15
0.3
0.2
Tran
0.1
6.2
Tmin
0.2
Model
0.5
0.5
0.2
Test
6.1
60
0.15
0.10
Test
Top
0.1
6.2
0.4
Test
6.1
(15.53) (100, >502
=500
(100
>500
Figure 10: Additional metrics for the UniRei50 results. We show the sequence separation range as columns. On each row, vary over cutoff thresholds L, L/5 and AUC defined in Appendix C.2. The top 3 rows are for train sequences, bottom 3 rows for test sequences.
19


Reviews:
[{"review_id": "o4uDVmiPlb-", "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": "This paper aims to improve low-depth MSAs, when a protein of interest only has a small number of known evolutionarily related sequences. This is a well motivated problem. MSAs are commonly used for a variety of purposes. Methods to enhance low-depth MSAs can be very useful. In particular, this paper focuses on using MSAs for contact prediction as a down-stream task. I'm not sure if contact prediction is the best use case for this, but it's a well-studied task for proof of concept.\n\nThe weight-sharing Neural Potts Model approach is novel (as far as I know) and is a well-motivated idea. There could be reasons to believe that weight sharing could improve Potts Model compared to directly fitting Potts Model to individual MSAs. Potts Model is very commonly used and it makes sense to base their method on Potts Model.\n\nPersonally, I think this approach is interesting, but the paper needs further experimental validation for acceptance. There are two main limitations in the experiments:\n\n   1. The paper claims that there is a gain from amortization for low-depth MSAs. This claim is only partially supported for very shallow MSAs. According to the experiments in the paper, there is only a gain for very shallow MSAs (fewer than 15 sequences) for long-range contact prediction, or for MSAs with fewer than 50 sequences for medium-range contact prediction. I'm not sure whether it's a common situation to have only 15 sequences in an MSA. It would be helpful if the authors could provide example use cases where this applies.\n\n   2.  The UniRef50 clusters are based on 50% sequence identity. In the paper, the train-test split is based on the clusters. There could still be significant information leakage if some train and test sequences share 40% similarity. I would like to see the authors further examine this and show that this is not the case.\n\nSuggestions for improving the paper:\n\n  1.  Include more than one PFAM clan level study. The P-loop NTPase example is well taken. It seems like the results here are quite different from the UniRef50 results (NPM has advantage for up to MSA depth 100s vs 15). This could be explained by the families in the same clan are a lot more similar than in UniRef50 in general. It would be helpful to understand why the NPM approach worked particularly well on NTPase, and whether it works similarly well on other PFAM clans. It seems like instead of aiming to use the method on arbitrary sequences, the method could be more valuable in cases where we have MSAs for other related families in the same clan.\n\n  2.  Split train-test differently in UniRef50 to make sure similar families don't end up in both train and test.\n\n  3. Instead of only evaluating the top L/5 contact prediction task, try other thresholds. The achieved contact prediction precision is around 10-20%. That seems rather low for both NPM and direct Potts, so hard to say if the gain is meaningful.\n\n  4. (Perhaps out of scope for this paper.) Could this method be applied to generate Potts Model parameters to be fed into other structure prediction models (e.g. AlphaFold). For example, it would be interesting to see that the NPM Potts Model can improve AlphaFold contact prediction results. \n\n  5. References for important proteins with shallow MSAs.\n\nClarification questions:\n\n1. How do you compute the MSA? Is it from HHblits or some other standard package?\n2. In the PFAM experiment, how do you vary the MSA Depth? Is it random subsampling?\n\n--------------------------------------------------------------------------------------------------------------\nUpdate:\nUpdated score from 5 to 6. The author response clarified some key questions and the updated paper incorporated some of the feedback.", "extra_scores": {}, "extra_text": {}}, {"review_id": "jGK-5l3MqJ0", "reviewer": "AnonReviewer2", "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": "The paper proposes a new object called Neural Potts Model (NPM) to train a Transformer to learn the local energy landscape of protein sequences. The problem of modeling energy landscapes using the power of techniques in natural language processing (NLP) is a timely and interesting problem. However, there are some concerns that limit the strength and the main claim of the paper that needs to be addressed. \n\nIs the NPM objective in equation (7) derived from some probabilistic model? While the intuition behind the objective can sound plausible from the discussion that follows, I am still curious what problem formulation is this objective solving? Is it a proper likelihood term?\n\nSince the paper is advocating the use of NPM objectives, is it possible to plot the value of the proposed NPM objective in the Transformers as a function of the Top-L precision? This will strengthen the claim in proving all the gains are actually coming from the objective and not other possible factors in training the model.\n\nThe paper can be improved by more effort in punctuation and proper labeling of figures. There are many sentences and phrases that need a comma for better readability. There are extra parentheses when referencing Figures. What are the axes showing in Figure 2? What are the highlighted shades indicating Figure 3 (std or sem)? Is it possible to have error bars for Figure 4? In abstract, MSA is not defined.\n\nAlso given that there is sufficient space left in the paper, some material including the Algorithm box can be moved to the main text. Also, the author can consider to use the space to expand on the significance and importance of the NPM objective.\n\n** after rebuttal: thanks for addressing the comments. I have revised my score based on the discussion.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "RR-bm_lhDG-", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "###########################################################################################################\n\n\nSummary of Paper\nThe motivation for the paper is a bit unclear. In the introduction, the authors begin by claiming to extend self-supervision \"to information from a set of evolutionarily related sequences\". However, it does not appear that the model is at all used for pretraining / representation learning as would be expected of a self-supervision method. No further connections to self-supervision are made in the rest of the paper. Based on the conclusion it seems that self-supervision is a future direction. If so, I would suggest rewriting the introduction to more clearly state the motivation of this paper.\n\nThe rest of the paper is more clear. The authors tackle the problem of predicting protein contacts. Standard unsupervised approaches fit a Potts model (often with pseudolikelihood) and then use the parameters of the fit model to make predictions about contacts. Such approaches require a MSA of reasonable depth and the amortized optimization approach suggested by the authors has the potential to work better than standard unsupervised approaches in the small-MSA regime by sharing information from all the protein families that the model is fit on.\n\nThe authors propose a meta-learning approach wherein a neural network (here a transformer) takes as input a single sequence and outputs the parameters of a Potts model. The model is trained with a pseudolikelihood loss across all sequences within individual families. This is a new and clever idea.\n\nThe main experimental result is that the NPM outperforms standard Potts models for shallow MSAs. This demonstrates that there apears to be some utility to this approach. However, the precision is still very poor and the authors make no claim that the enhanced accuracy is of biological utility (i.e. can be used to fold the protein). There are other tools that perform contact prediction from a single sequence with higher accuracy than the NPM. Furthermore NPM performs worse than standard Potts models for medium and large MSAs. This seems to be a significant drawback of the method. \n\n#############################################################################################\n\n  \nQuestions / Suggested Experiments:\n  \nPFAM EXPERIMENT: \n\nCan the authors please explain why the PFAM experiments are evaluated on top-L precision whereas the remaining experiments are top-L/5 precision? The lack of explanation in the paper suggests the metrics may be cherry-picked to best show the performance of the model. \n\nSince all the proteins in the family share structural similarity, it is not clear that NPM has learned to transfer any information. After looking at the structures in PyMOL it seems that significant substructure is shared between the different families. Here is a resonable baseline/experiment to add: Align the query sequence to the higher-depth MSAs and then using the MSA that best fits the query sequence, predict contacts based on the individual Potts model trained on that family. This would elucidate whether or not NPM is simply treating the query sequence as if it were from one of the higher-depth MSAs.\n\nUNIREF EXPERIMENT: \n\nThe details of this experiment could use more explanation. I did not see any explanation of the purpose of the training and heldout sets, which in this setting are not obvious. My guess is that the authors are trying to demonstrate that the NPM can generalize zero-shot to new families? If so, please mention it in the text. Random splits are problematic for demonstrating zero-shot generalization since there could be very similar families in the train and heldout sets (e.g. same PFAM clan or structural class). As written, the paper barely discusses the purpose of the heldout set but if it is meant to convey some test of generalization the authors should more carefully construct their heldout test set accordingly.\n\nWhy is there no plot for short-range predictions, as there was for the PFAM experiment?\n\nThere is an unclear statement: \"During training, we iterate over all sequences and their MSAs on every epoch, and subsample to M=30 sequences per MSA.\" Is this just for NPM or also for the Potts model? I would be very concerned if this is also what is done for the Potts models as it would significantly harm their performance. The authors need to clarify this in the text of the paper.\n\nThe authors do not discuss other existing methods that currently exists for shallow MSAs. For example, standard semi-supervised methods are fine-tuned on contact prediction and thus, work with single-sequence inputs. I believe these methods may perform as well or better than NPM. The authors should discuss these methods and include comparisons. \n\n\nModel Architecture\nPlease do the following:\n1) In Appendix B.1 you describe a number of \"tricks\" to reduce the number of parameters. These include (1) a low-rank decomposition of the bilinear form, (2) weight tying by amino acid for the decompositions (3) Convolutional layers. Please provide experiments showing the utility (or lack thereof) of each of these. This will lift the results from \"here is what happens if you use these tricks\" to \"here is data suggesting that these choices we made are better than the naive choices\". Thus, add more impact to the paper.  \n\n2) No explanation is given for why the architectures are so different for PFAM and UniRef (e.g. convolutional layers for PFAM only). Please provide an explanation. A priori, I see no reason to use different architectures for these two different domains. It seems to only complicate the paper and methods. If these choices really are important, then this seems to suggest overfitting to particular tasks rather than having a general solution.\n\n\nAppendix Training Details:\nI personally found the written descriptions here to be too vague to understand the exact details. This made it harder to evaluate the paper. Can you please clarify:\n1) \"we randomly subsample the MSA down to 100 sequences as NPM target sequences...\". Based on equation (7) this means the pseudo likelihood is evaluate on 100 random sequences from the MSA. This seems to directly contradict the experiment shown in Figure 3 where the number of sequences in the MSA is varied up to 1000. Please explain what is happening here?\n\n\n###########################################################################################\n\nExplanation of Score\n\nThis paper was difficult to score because I find the idea / approach to be very creative and different from standard techniques of borrowing the latest NLP pertaining task and applying it to proteins. I applaud the authors for taking a unique approach. \n\nThe main drawback is that the authors exclusively evaluate the model on contact prediction and do not demonstrate convincing performance. While the model sometimes outperforms standard approaches with shallow MSAs the performance is still quite bad and no comparisons to other methods are shown. \n\nThus, overall I am impressed with the direction of the paper but think it needs more work. \n\n\nUpdated score from 5 -> 6 after clarifying feedback from the authors. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "9u-kxZKIe6k", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "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": "Overall comments\n- The paper introduces an amortized optimization framework to learn a mapping between an input sequence and the coefficients of the corresponding Potts model -- a standard method used in computational biology to model protein sequences.\n- The idea is particularly compelling in the regime where there exists a limited number of sequences (in the training data) that are similar to the target sequence. That would typically result in a low quality multiple sequence alignment (MSA), and in turn a poor Potts model.\n- The idea is sound from a theoretical standpoint. The newly introduced framework is well motivated: authors did a good job at explaining the background for the problem and providing an intuition for why the proposed amortized optimization framework would improve over the current approach (i.e. independently training models on each MSA).\n- The paper is very well written and easy to read -- the language is precise, everything is defined very clearly. \n- Experiments show good results in the regime that matters most for the introduced Neural Potts model (NPM), i.e., low density of sequences around the query.\n- The approach would be even more compelling if the NPM would not underperform the baseline in the regime where Meff (the “effective number” of sequences in the MSA) is larger. \n- Bridging that gap would be very compelling from a practical standpoint: large scale studies typically involve modeling thousands of distinct proteins, and thus fitting thousands of distinct models independently. Intuitively, there is shared information between these models, which is lost when models are trained independently but could be captured via a process akin to the one introduced in this paper.\n\nDetailed comments and questions\n\nSection 1\n- Would suggest to add one sentence describing why it matters to solve the low Meff regime to further motivate your approach (e.g., which scientific questions would we be able to better answer?).\n\nSection 2\n- I would add a paragraph on Transformers models for protein sequence embedding (since it is a core component of the NPM used in experiments), and cite a few of the key works in this area, for example:\n- Rives et al., Biological Structure and Function Emerge from Scaling Unsupervised Learning to 250 Million Protein Sequences (https://www.biorxiv.org/content/10.1101/622803v1)\n- Madani et al., ProGen: Language Modeling for Protein Generation (https://www.biorxiv.org/content/10.1101/2020.03.07.982272v2)\n- Vig et al., Bertology meets biology: Interpreting attention in protein language models (https://arxiv.org/abs/2006.15222)\n- Rao et al., Evaluating Protein Transfer Learning with TAPE (https://arxiv.org/abs/1906.08230)\n\nSection 3\n- Is there a particular reason for which you are using a vanilla Transformer architecture? Other Transformer architectures (e.g., BERT-like architecture) may be able to learn better sequence embeddings and in turn further close the gap with the independent Potts model in the large Meff regimes (see the above papers for example).\n\nSection 4\n- How do you interpret the gap with the baseline Potts model approach in the medium-large Meff setting? \n- Is there a way to detect (based on training) the situations where the NPM will likely underperform the independent Potts model baseline?\n- Computationally, how expensive is it to train one NPM Vs one Potts model? For example, if I have 1000 proteins to model in the regime where both perform the same (Meff ~100), is it faster for me to train one NPM model or 1000 Potts models? I suspect the 1000 Potts models are still cheaper to train, but if not that could be one strength of your approach to mention.\n- What is the distribution of Meff over the data (section 4.2)? Based on the rightmost plots it seems that the majority of MSAs are in the “>500” regimes where NPM underperforms independent Potts models (asking that question with the understanding that you are more interested by addressing the low Meff setting here).\n- Do you have an intuition for why the NPM appears to be doing better for medium range interaction here Vs long range in experiment 4.1?\n- How do you explain the U-shape of the top-L/5 precision curves for NPM (on heldout data)? I would have expected that increasingly larger Meff would be associated with monotonically increasing performance (as is the case for the Independent Potts model curves in blue).\n\nAppendix C1\n- Why the upweight by a factor sqrt(Meff(n))?\n\nAppendix C3\n- Did you look at the performance gap on HTH? Since baseline Potts models yield poor long-range contact on that dataset, perhaps NPM would have helped address that issue, given increased ability to model long range dependencies as per 4.1. It was not clear whether the low performance is due to an intrinsic limitation of the Potts model or potentially due to a low density around each query point in HTH (in which case NPM may help).\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This is a creative piece of work wherein learning of what is normally family-specific Potts models is turned into an amortized optimization problem across different families of proteins. The Potts models are learned with a pseudolikelihood approach, and the evaluation of the model against baselines is performed only on a contact prediction problem. This last point is problematic, because on the one hand, the authors use this "as a proxy for the underlying accuracy of the Potts model learned", and on the other hand, claim that "we do not want to claim our method is state-of-the-art for contact prediction, it is certainly not".  Overall, the paper is promising, but is too preliminary on the empirics to warrant publication at this time.
Decision: Reject