Neural Potts Model

1 INTRODUCTION . 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 . 2 BACKGROUND . 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 . 2.1 AMORTIZED OPTIMIZATION . 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 ) . 3 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∗ } = 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 can not 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 .

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