Year: 2020

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.

heading: 1 INTRODUCTION
text: When two positions in a protein sequence are in spatial contact in the folded three-dimensional structure of the protein, evolution is not free to 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 a Potts model energy function to the MSA (Lapedes et al., 1999; Thomas et al., 2008; Weigt et al., 2009). Various approximations are used in practice including mean field (Morcos et al., 2011), sparse inverse covariance estimation (Jones et al., 2011), and pseudolikelihood maximization (Balakrishnan et al., 2011; Ekeberg et al., 2013; Kamisetty et al., 2013). To construct the MSA for a given input sequence, a 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 on 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 is 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 a single model with unsupervised learning across many evolutionarily diverse protein sequences. We extend this core idea to train a model to output an 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., 2011; Kamisetty et al., 2013; Moult et 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 an independent Potts model to each family.
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 a 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.

heading: 2 BACKGROUND
text: Multiple sequence alignments An MSA is a set of aligned protein sequences that are evolutionarily related. MSAs are constructed by retrieving related sequences from a sequence database and aligning the returned sequences using a heuristic. An MSA can be 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 (x1, . . . , xL) of length L with the negative energy function:
−E(x) = ∑ i hi(xi) + ∑ ij Jij(xi, xj) (1)
Which defines potentials hi 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(x;W ) is obtained by normalization as exp{−E(x;W )}/Z(W ). Since the normalization constant is intractable, pseudolikelihood is commonly used to fit the parameters (Balakrishnan et al., 2011; Ekeberg et al., 2013). Pseudolikelihood approximates the likelihood of a sequence x as a product of conditional distributions: `PL(x;W ) = − ∑ i log p(xi|x−i; W ). To estimate the Potts model, we take the expectation:
LPL(W ) = E x∼M [`PL(x;W )] (2)
over an MSAM. In practice, we have a finite set of sequences M̂ in the MSA to estimate Eq. (2). L2 regularization ρ(W ) = λJ ‖J‖2 + λh ‖h‖2 is added, and sequences are reweighted to account for redundancy (Morcos et al., 2011). We write the regularized finite sample estimator as:
L̂PL(W ) = 1
Meff M∑ m=1 wm[`PL(x m;W )] + ρ(W ) (3)
Which sums over all the M sequences of the finite MSA M̂, weighted with wm summing collectively to Meff . The finite sample estimate of the parameters Ŵ ∗ is obtained by minimizing L̂PL.
Idealized MSA Notice how in Eq. (2), we idealized the MSAM 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 a hypothetical idealized estimate of the parameters to explain how the Neural Potts Model can improve on the finite sample estimate Ŵ ∗ for low-depth MSAs.

heading: 2.1 AMORTIZED OPTIMIZATION
text: We review amortized optimization (Shu, 2017), a 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
W ∗(x) = argmin W L(W ;x), (4)
where W ∈ Rm is the optimization variable, x ∈ Rn is the input or conditioning variable to the optimization problem, and L : Rm × Rn → R is the objective. We assume W ∗(x) is unique. We consider the setting of having a distribution over optimization problems with inputs x ∼ p(x), and the argmin of those optimization problems W ∗(x).
Amortization uses learning to leverage the shared structure present across the distribution, e.g. a solutionW ∗(x) is likely correlated with another solutionW ∗(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 Wθ(x) such that hopefully Wθ ≈W ∗. Modeling and learning Wθ(x) are the key design decisions when using amortization.
Modeling approaches. In this paper we consider models Wθ(x) that directly predict the solution to Eq. (4) with a neural network, an 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 L(W ;x) and gradient information ∇WL(W ;x), e.g. by predicting multiple candidate solutions W and selecting the most optimal one. This is sometimes referred to as semi-amortization or unrolled optimization-based models and is considered in Gregor & LeCun (2010) for sparse coding, Li & Malik (2016); Andrychowicz et al. (2016); Finn et al. (2017) for meta-learning, and Marino et al. (2018); Kim et al. (2018) for posterior optimization.
Learning approaches. There are two main classes of learning approaches for amortization:
argmin θ E p(x) L(Wθ(x);x) (5) argmin θ E p(x) ‖Wθ(x)−W ∗(x)‖22 . (6)
Gradient-based approaches leverage gradient information of the objective L and optimize Eq. (5) whereas regression-based approaches optimize a distance to ground-truth solutions W ∗, such as the squared L2 distance in Eq. (6). Prior work has shown that models trained with these objectives can learn to predict the optimal W ∗ directly as a function of x. 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 solutionW ∗ while regression-based approaches are 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 et 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 et al., 2007; Bain & Sammut, 1995).

heading: 3 NEURAL POTTS MODEL
text: In Eq. (2) we introduced the Potts model for a single MSA M (aligned set of sequences x), to optimize W ∗ = {h∗, J∗} = argminW Ex̃∼M[`PL(x̃;W )]. As per Eq. (5) We will now introduce a neural network to estimate Potts model parameters from a single sequence: {hθ(x), Jθ(x)} =Wθ(x) with a single forward pass.
We propose minimizing the following objective for the NPM parameters θ, which directly minimizes the Potts model losses in expectation over our data distribution x ∼ D and their MSAs x̃ ∼M(x):
LNPM(θ) = E x∼D
[ E
x̃∼M(x) `PL(x̃;Wθ(x))
] (7)
To compute the loss for a given sequence x we compute the Potts model parameters Wθ(x), and evaluate its pseudo-likelihood loss `PL on a set of sequences x̃ from the MSA constructed with x as query sequence. This fits exactly in “amortized optimization” in Section 2.1 Eq. (5): we train a model to predict the outcome of a set of highly related optimization problems. One key extension to the described amortized optimization setup is that the model Wθ estimates the Potts Model parameters from only the MSA query sequence x as input rather than the full MSAM(x). Thus, our model must learn to distill the protein energy landscape into its parameters, since it cannot look up related proteins during runtime. A full algorithm is given in Appendix A.
Similar to the original Potts model, we need to add a regularization penalty ρ(W ) to the main objective. For a finite sample of N different query sequences {xn}, and a corresponding sample of N× M aligned sequences {x̃mn } from MSA M̂(xn), the finite sample regularized loss, i.e. NPM training objective, becomes:
L̂NPM(θ) = N∑ n=1
[ 1
Meff(n) M∑ m=1 wmn [`PL(x̃ m n ;Wθ(xn))] + ρ(Wθ(xn))
] (8)
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 Wθ 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 Wθ(x) will be behind the optimal W ∗ for the objective L: L(Wθ(x)) > L(W ∗). This is closely related to underfitting: the model Wθ is not flexible enough to capture W ∗(x). However, recall that in the Potts model setting, there is a finite-sample training objective L̂ (Eq. (8)), with minimizer Ŵ ∗. We can expect an amortization gap in the training objective; however this amortization gap can now be advantageous. Even if the amortized solution Wθ(x) is near-optimal on L̂, it can likely find a more generalizable region of the overparametrized domain W by parameter sharing of θ, allowing it to transfer information between related instances. The inductive bias of Wθ(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 Wθ, its pre-training, as well as the shared structure between the protein families in the dataset. Concretely we will consider for the generalization loss L not just the pseudo-likelihood 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.
We will show that for some samples L(Wθ(x)) < L(Ŵ ∗), i.e. there is an inductive generalization gain. This is visually represented in Fig. 2; and Table 1 compares amortized optimization and NPM, making a connection to multi-task learning (Caruana, 1998). Additionally, we could frame NPM as a hypernetwork, a neural network that predicts the weights of second network (in this case the Potts model) as in, e.g., Gomez & Schmidhuber (2005); Ha et al. (2016); Bertinetto et al. (2016).
In summary, the goal for the NPM is to “distill” an ensemble of Potts models into a single feedforward model. From a self-supervised learning perspective, rather than supervising the model with the input directly, we use supervision from an energy landscape around the input.

heading: 4 EXPERIMENTS
text: In Section 4.1 we present results on a small set of related protein domain families from Pfam, where we artificially reduce the MSA depth for a few families to study the inductive generalization gain from the shared parameters. In Section 4.2 we present results on a large Transformer trained on MSAs for all of UniRef50.
For the main representation gθ(x) we use a bidirectional transformer model (Vaswani et al., 2017). To compute the four-dimensional pairwise coupling tensor Jθ(x) from sequence embedding gθ(x) we introduce the multi-head bilinear form (mhbf) in Appendix B. One can think of the multi-head bilinear form as the L× L self-attention maps of the Transformer’s multi-head attention module, but without softmax normalization. When using mhbf for direct prediction, there are K2 heads, one for every amino acid pair k, l. 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 gθ(x) with a Transformer pre-trained with masked language modeling following (Rives et al., 2019).
To evaluate Neural Potts Model energy landscapes, we will focus on 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.

heading: 4.1 PFAM CLANS
text: To study generalization in a controlled setting, we investigate a small set of structurally-related MSAs from the Pfam domain family database (Finn et al., 2016) belonging to 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 a Pfam clan are linked by a distant evolutionary relationship, giving them related but not trivially-similar structure. We obtain contact maps for the sequences in each of the families where a structure is available in the PDB. At test time we input the sequence and compare the generated couplings under the model to the corresponding structure.
We compare the NPM to two baselines. The first direct comparison is to an independent Potts model trained directly on the MSA. For the second baseline we construct the “nearest neighbor” Potts model, by aligning each test 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 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 a 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 on 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 on the reduced MSA but evaluated using only the target sequence as input. We train a separate NPM on each of the five cross-evaluation rounds, evaluate on the structures corresponding to the bucket with reduced
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 on 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 a 12-layer Transformer with protein language modeling pre-training. Because of the small dataset size, we keep the weights of the base Transformer gθ 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 a higher precision relative to 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 is 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 pseudo-likelihood loss (Eq. (7)) on the MSAs 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 we see expected overfitting for smaller MSA depth: better training loss but worse contact precision.
Additionally, we 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 are 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
Range = Medium
Range = Long
Top-L/5,LR precisions
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 et 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).
We ask whether an amortization gain can 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 a sample from the MSA of the input sequence. In each training epoch, we randomly subsample a different set of 30 sequences from the MSA to fit the NPM objective. We use ground-truth structures to 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.1.
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 a 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 in 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 us 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 a single feedforward pass without access to 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. it 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 is trained on UniRef50.

heading: 5 RELATED WORK
text: Recently, protein language modeling has emerged as a 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 et al. (2020) trained models with autoregressive objectives. Transformer protein
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 on 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 a single sequence; the second is to use amortized optimization to 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 a supervised pipeline.
Unsupervised structure learning is reviewed in the introduction. The main approach has been to learn a set of constraints from a 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 on 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 et 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 a logical extension of this work to integrate the Neural Potts model into the supervised pipeline.

heading: 6 DISCUSSION
text: This paper explores how a protein sequence model can be trained to produce a local energy landscape that is defined by a set of evolutionarily related sequences for each input. The training objective is cast as an amortized optimization problem. By learning to output the parameters for a 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 a single model using the amortized objective to 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 architecture and parameterization of the energy function by the deep network, use of the amortized models in a 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 is to directly incorporate information from a 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.

heading: Appendix

heading: A LEARNING THE NEURAL POTTS MODEL
text: Algorithm 1 Learning the Neural Potts Model Initialize the amortized model Wθ while unconverged do
Sample N sequences xn ∼ D and M MSA sequences x̃mn ∼M(xn) grad_update(θ,∇θL̂NPM(θ; {xn}, {x̃mn }) . Optimize the regularized loss in Eq. (8)
end while

heading: B MODEL ARCHITECTURE: MULTI-HEAD BILINEAR FORM FOR PAIRWISE COUPLINGS
text: In this Section, we describe the model architecture to compute a four-dimensional pairwise coupling tensor Jθ(x) from sequence embedding gθ(x).

heading: B.1 MULTI-HEAD BILINEAR FORM
text: We write sequence lengthL and amino acid vocabularyK = 21. The single site potentials h ∈ RL×K , and the pairwise couplings are a four-dimensional tensor: J ∈ RL×K×L×K indexed as Jij(k, l). We start with a sequence-level model to produce the embedding e of the sequence (typically final hidden layer output): e = gθ(x) ∈ RL×d. The estimator for single-site potential hθ(x) is a linear projection layer on the embedding; hθ(x) = gθ(x)Ph with Ph ∈ Rd×K . Now we discuss how to parametrize the estimator Jθ(x) ∈ RL×K×L×K .
Multi-head bilinear form for direct prediction We introduce a multi-head bilinear form (mhbf) on the embedding e; i.e. for every pair k, l of amino acids we have a bilinear form, parametrized with a learned interaction matrix Bkl ∈ Rd×d connecting the hidden states at positions ei, ej ∈ R1×d. So we compute the K2 bilinear forms for amino acid pairs (k, l) between L × L position pairs (i, j): Jij(k, l) = eiBkle>j . We always use a low-rank decomposition B
kl = UklV kl> with both Ukl, V kl ∈ Rd×d′ , so the bilinear form becomes the inner product in the lower-dimensional space with d′ the projection dimension: (eiUkl)(ejV kl)>. We can interpret this as an inner product of embeddings i, j after linear projection to a space specific to amino acid pair (k, l). 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 θ include the parameters of the transformer that produces the embedding and the components of the interaction matrix {Ukl, V kl}.
Direct prediction: tied projection One way to reduce the number of parameters in the multi-head bilinear form, is for the low-rank decompositions of the K2 heads Bkl to share their decomposition per k, l. Instead we can share/tie the projection matrices between amino acids k, l: Ukl = Uk and V kl = V l, such that head Bkl = UkV l>. Note that the dot product in this case is after a linear projection specific to single-site amino acid k and l separately; Jij(k, l) = eiBkle>j = (eiU k)(ejV l)>.
Direct prediction: Symmetry We can or should parametrize the estimator of J to be symmetrical against interchanging both i,j and k,l: Jij(k, l) = Jji(l, k), i.e. no difference between the order of considering interaction between AA k at position i with AA l at position j. This does not mean symmetry of each interaction matrix! We ask that
Jij(k, l) = eiB kle>j = ejB lke>i = eiB lk>ej
The second equality is the symmetry, the last equality by transposing the bilinear form. From Bkl = Blk> it follows that
UklV kl> = V lkU lk>
Ukl = V lk
The last equality is the obvious choice. In the tied parametrization, this simply becomes Ukl = Uk = V k = V lk such that Wkl = UkU l>. Once again, note that the dot product now becomes Jij(k, l) = (eiU k)(ejU l)>. We present a Tensor decomposition perspective on this multi-head bilinear form in Appendix B.2.
Convolutional layers after multi-head bilinear form As an extended model architecture, we consider having convolutional layers after the multi-head bilinear form (only used for the Pfam experiments). parametrized with a learned interaction matrix Bkl ∈ Rd×d In this case, rather than having K2 heads Bkl, we now have an arbitrary number of heads F which will become the number of channels in the consecutive convolutional layers: BF = UFV F>. We add 1× 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 in this model variation.

heading: B.2 TENSOR DECOMPOSITION VIEW ON MULTI-HEAD BILINEAR FORM
text: We can see the multi-head bilinear form as a tensor decomposition of J , for which we 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 Ukl matrices as U ∈ RK×K×d×d′ ; and index into U in the same notation as for U : Uklαr = Uklαr. With α, β ∈ [1 . . . d], r ∈ [1 . . . d′] The same for V . Now the J estimate in the full untied asymmetric case, written as tensor, becomes
Jklij = e α i Uklαr Vklrβ eβj
or the symmetric (Uklαr = V lkαr) and tied (U ∈ RK×d×d ′ ) version:
Jklij = e α i Ukαr U lrβ eβj
Note that the U ,V are shared across proteins, while the embeddings e = gθ(σ) are specific per protein, based on a high-capacity sequence level model.

heading: C EXPERIMENT DETAILS

heading: C.1 TRAINING DETAILS
text: We summarize the precise model architecture and optimization settings in Table 2. During each NPM training step, for a given input x, M sequences x̃m 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 M sampled sequences as similar to a minibatch. Note that to 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 √ Meff(n), which places more weight on the well-formed, deep MSAs and discounts the shallower MSAs. In both the Pfam and UniRef experiments, we enforce a max sequence length of 512 via random contiguous crops of positions.
C.2 VALIDATION DETAILS
To compute precisions, we convert the pairwise couplings J ∈ RL×K×L×K to an L× 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 ×K matrix Jij for every pair of positions i, j, and (3) applying the Average Product Correction (Dunn et al., 2008). True contacts are defined as pairwise distances less than or equal to 8 Angstroms. Precision is calculated as the true positive fraction of the top L, bL/2c, or bL/5c predicted contacts. Additional to precision, the Area Under the
Precision-Recall Curve (AUC) is computed, summing over thresholds stepwise per L/10 increment up to L. Precision and AUC metrics are computed at sequence separations s of short (6 ≤ s < 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), a 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.

heading: C.3 PFAM TRAINING DATA AND SETUP
text: 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 a 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 a five-fold cross-evaluation scheme, in which we partition the clan’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, 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 a 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 in the four buckets that have not been reduced, as well as the reduced bucket. At a given training step, we randomly select a sequence x 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 x̃ from the MSA to fit the NPM objective. The procedure is described in more detail in Appendix C.1.
Evaluation During evaluation, we assess the NPM and the independent Potts model via a contact prediction task (described in previous subsection), on the families in the evaluation bucket. For
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 a 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 Pfam seed alignment as input. We select the family with the highest HHalign probability score as the nearest neighbor. The nearest neighbor prediction is generated as follows: (1) the validation sequence is aligned to the selected train family’s MSA; (2) an independent Potts model is fit to the selected train family’s MSA, yielding a predicted contact map for the train family; (3) the rows and columns of the predicted contact map that align to the validation sequence are extracted to construct a prediction for the validation sequence.

heading: C.4 UNIREF50 TRAINING DATA AND SETUP
text: For the experiments in Section 4.2, we retrieve the UniRef-50 (Suzek et al., 2007) database dated 2018- 03. The UniRef50 clusters are partitioned randomly in 90% train and 10% test sets. For all sequences, we construct MSAs using HHblits (Steinegger et al., 2019) against the UniClust30_2017_10 database. HHblits is run using the default settings, for 3 iterations with an e-value of 0.001.
It is important to note that given this MSA generation procedure, validation sequences can be included in MSAs of train sequences. However, we are guaranteed that validation sequences are not trained on as input to NPM.

heading: Evaluation of contact precision
text: 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, we 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 on sequences from its training set, and sequences from its test set. Note, that
the structures are 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 a list of all protein structures with a resolution less than 2.5 Å, a 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. If the PDB entry retrieves hits only to training sequences we assign it to the training-sequences group. If the PDB entry retrieves hits only to test sequences we assign it to the test-sequences group. Any PDB entry which hits both training and test sequences or neither, is discarded. To perform the search we 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 at 50% sequence similarity, resulting in a dataset of 11040 structures assigned to train sequences and 211 structures assigned to 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.
D ADDITIONAL EXPERIMENTS

Review summaries:
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. 

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.