Transformer protein language models are unsupervised structure learners

1 INTRODUCTION . Unsupervised modeling of protein contacts has an important role in computational protein design ( Russ et al. , 2020 ; Tian et al. , 2018 ; Blazejewski et al. , 2019 ) and is a central element of all current state-of-the-art structure prediction methods ( Wang et al. , 2017 ; Senior et al. , 2020 ; Yang et al. , 2019 ) . The standard bioinformatics pipeline for unsupervised contact prediction includes multiple components with specialized tools and databases that have been developed and optimized over decades . In this work we propose replacing the current multi-stage pipeline with a single forward pass of a pre-trained end-to-end protein language model . In the last year , protein language modeling with an unsupervised training objective has been investigated by multiple groups ( Rives et al. , 2019 ; Alley et al. , 2019 ; Heinzinger et al. , 2019 ; Rao et al. , 2019 ; Madani et al. , 2020 ) . The longstanding practice in bioinformatics has been to fit linear models on focused sets of evolutionarily related and aligned sequences ; by contrast , protein language modeling trains nonlinear deep neural networks on large databases of evolutionarily diverse and unaligned sequences . High capacity protein language models have been shown to learn underlying intrinsic properties of proteins such as structure and function from sequence data ( Rives et al. , 2019 ) . A line of work in this emerging field proposes the Transformer for protein language modeling ( Rives et al. , 2019 ; Rao et al. , 2019 ) . Originally developed in the NLP community to represent long range context , the main innovation of the Transformer model is its use of self-attention ( Vaswani et al. , 2017 ) . Self-attention has particular relevance for the modeling of protein sequences . Unlike convolutional or recurrent models , the Transformer constructs a pairwise interaction map between all positions in the sequence . In principle this mechanism has an ideal form to model protein contacts . In theory , end-to-end learning with a language model has advantages over the bioinformatics pipeline : ( i ) it replaces the expensive query , alignment , and training steps with a single forward ∗Work performed during an internship at Facebook . 1Weights for all ESM-1 and ESM-1b models , as well as regressions trained on these models can be found at https : //github.com/facebookresearch/esm . pass , greatly accelerating feature extraction ; and ( ii ) it shares parameters for all protein families , enabling generalization by capturing commonality across millions of evolutionarily diverse and unrelated sequences . We demonstrate that Transformer protein language models learn contacts in the self-attention maps with state-of-the-art performance . We compare ESM-1b ( Rives et al. , 2020 ) , a large-scale ( 650M parameters ) Transformer model trained on UniRef50 ( Suzek et al. , 2007 ) to the Gremlin ( Kamisetty et al. , 2013 ) pipeline which implements a log linear model trained with pseudolikelihood ( Balakrishnan et al. , 2011 ; Ekeberg et al. , 2013 ) . Contacts can be extracted from the attention maps of the Transformer model by a sparse linear combination of attention heads identified by logistic regression . ESM-1b model contacts have higher precision than Gremlin contacts . When ESM and Gremlin are compared with access to the same set of sequences the precision gain from the protein language model is significant ; the advantage holds on average even when Gremlin is given access to an optimized set of multiple sequence alignments incorporating metagenomics data . We find a linear relationship between language modeling perplexity and contact precision . We also find evidence for the value of parameter sharing : the ESM-1b model significantly outperforms Gremlin on proteins with low-depth MSAs . Finally we explore the Transformer language model ’ s ability to generate sequences and show that generated sequences preserve contact information . 2 BACKGROUND . Multiple Sequence Alignments ( MSAs ) A multiple sequence alignment consists of a set of evolutionarily related protein sequences . Since real protein sequences are likely to have insertions , deletions , and substitutions , the sequences are aligned by minimizing a Levenshtein distance-like metric over all the sequences . In practice heuristic alignment schemes are used . Tools like Jackhmmer and HHblits can increase the number and diversity of sequences returned by iteratively performing the search and alignment steps ( Johnson et al. , 2010 ; Remmert et al. , 2012 ) . Metrics For a protein of length L , we evaluate the precision of the top L , L/2 , and L/5 contacts for short range ( |i − j| ∈ [ 6 , 12 ) ) , medium range ( |i − j| ∈ [ 12 , 24 ) ) , and long range ( |i = j| ∈ [ 24 , ∞ ) ) contacts . We also separately evaluate local contacts ( |i−j| ∈ [ 3 , 6 ) ) for secondary structure prediction in Appendix A.9 . In general , all contacts provide information about protein structure and important interactions , with shorter-range contacts being useful for secondary and local structure , while longer range contacts are useful for determining global structure ( Taylor et al. , 2014 ) . 3 RELATED WORK . There is a long history of protein contact prediction ( Adhikari & Cheng , 2016 ) both from MSAs , and more recently , with protein language models . Supervised contact prediction Recently , supervised methods using deep learning have resulted in breakthrough results in supervised contact prediction ( Wang et al. , 2017 ; Jones & Kandathil , 2018 ; Yang et al. , 2019 ; Senior et al. , 2020 ; Adhikari & Elofsson , 2020 ) . State-of-the art methods use deep residual networks trained with supervision from many protein structures . Inputs are typically covariance statistics ( Jones & Kandathil , 2018 ; Adhikari & Elofsson , 2020 ) , or inferred coevolutionary parameters ( Wang et al. , 2017 ; Liu et al. , 2018 ; Senior et al. , 2020 ; Yang et al. , 2019 ) . Other recent work with deep learning uses sequences or evolutionary features as inputs ( AlQuraishi , 2018 ; Ingraham et al. , 2019 ) . Xu et al . ( 2020 ) demonstrates the incorporation of coevolutionary features is critical to performance of current state-of-the-art methods . Unsupervised contact prediction In contrast to supervised methods , unsupervised contact prediction models are trained on sequences without information from protein structures . In principle this allows them to take advantage of large sequence databases that include information from many sequences where no structural knowledge is available . The main approach has been to learn evolutionary constraints among a set of similar sequences by fitting a Markov Random Field ( Potts model ) to the underlying MSA , a technique known as Direct Coupling Analysis ( DCA ) . This was proposed by Lapedes et al . ( 1999 ) and reintroduced by Thomas et al . ( 2008 ) and Weigt et al . ( 2009 ) . Various methods have been developed to fit the underlying Markov Random Field , including meanfield DCA ( mfDCA ) ( Morcos et al. , 2011 ) , sparse inverse covariance ( PSICOV ) ( Jones et al. , 2011 ) and pseudolikelihood maximization ( Balakrishnan et al. , 2011 ; Ekeberg et al. , 2013 ; Seemayer et al. , 2014 ) . Pseudolikelihood maximization is generally considered state-of-the-art for unsupervised contact prediction and the Gremlin ( Balakrishnan et al. , 2011 ) implementation is used as the baseline throughout . We also provide mfDCA and PSICOV baselines . Recently deep learning methods have also been applied to fitting MSAs , and Riesselman et al . ( 2018 ) found evidence that factors learned by a VAE model may correlate with protein structure . Structure prediction from contacts While we do not perform structure prediction in this work , many methods have been proposed to extend contact prediction to structure prediction . For example , EVFold ( Marks et al. , 2011 ) and DCAFold ( Sulkowska et al. , 2012 ) predict co-evolving couplings using a Potts Model and then generate 3D conformations by directly folding an initial conformation with simulated annealing , using the predicted residue-residue contacts as constraints . Similarly , FragFold ( Kosciolek & Jones , 2014 ) and Rosetta ( Ovchinnikov et al. , 2016 ) incorporate constraints from a Potts Model into a fragment assembly based pipeline . Senior et al . ( 2019 ) , use features from a Potts model fit with pseudolikelihood maximization to predict pairwise distances with a deep residual network and optimize the final structure using Rosetta . All of these works build directly upon the unsupervised contact prediction pipeline . Contact prediction from protein language models Since the introduction of large scale language models for natural language processing ( Vaswani et al. , 2017 ; Devlin et al. , 2019 ) , there has been considerable interest in developing similar models for proteins ( Alley et al. , 2019 ; Rives et al. , 2019 ; Heinzinger et al. , 2019 ; Rao et al. , 2019 ; Elnaggar et al. , 2020 ; Lu et al. , 2020 ; Madani et al. , 2020 ; Shen et al. , 2021 ) . Rives et al . ( 2019 ) were the first to study protein Transformer language models , demonstrating that information about residue-residue contacts could be recovered from the learned representations by linear projections supervised with protein structures . Recently Vig et al . ( 2020 ) performed an extensive analysis of Transformer attention , identifying correspondences to biologically relevant features , and also found that different layers of the model are responsible for learning different features . In particular Vig et al . ( 2020 ) discovered a correlation between selfattention maps and contact patterns , suggesting they could be used for contact prediction . Prior work benchmarking contact prediction with protein language models has focused on the supervised problem . Bepler & Berger ( 2019 ) were the first to fine-tune an LSTM pretrained on protein sequences to fit contacts . Rao et al . ( 2019 ) and Rives et al . ( 2020 ) perform benchmarking of multiple protein language models using a deep residual network fit with supervised learning on top of pretrained language modeling features . In contrast to previous work on protein language models , we find that a state-of-the-art unsupervised contact predictor can be directly extracted from the Transformer self-attention maps . We perform a thorough analysis of the contact predictor , showing relationships between performance and MSA depth as well as language modeling perplexity . We also provide methods for improving performance using sequences from an MSA and for sampling sequences in a manner that preserves contacts . 4 MODELS . We compare Transformer models trained on large sequence databases to Potts Models trained on individual MSAs . While Transformers and Potts Models emerged in separate research communities , the two models share core similarities ( Wang & Cho , 2019 ) which we exploit here . Our main result is that just as Gremlin directly represents contacts via its pairwise component ( the weights ) , the Transformer also directly represents contacts via its pairwise component ( the self-attention ) . 4.1 OBJECTIVES . For a set of training sequences , X , Gremlin optimizes the following pseudolikelihood loss , where a single position is masked and predicted from its context . Inputs are aligned , so all have length L : LPLL ( X ; θ ) = E x∼X L∑ i=1 log p ( xi|xj 6=i ; θ ) ( 1 ) The masked language modeling ( MLM ) loss used by the Transformer models can be seen as a generalization of the Potts Model objective when written as follows : LMLM ( X ; θ ) = E x∼X E mask ∑ i∈mask log p ( xi|xj 6∈mask ; θ ) ( 2 ) In contrast to Gremlin , the MLM objective applied by protein language modeling is trained on unaligned sequences . The key distinction of MLM is to mask and predict multiple positions concurrently , instead of masking and predicting one at a time . This enables the model to scale beyond individual MSAs to massive sequence datasets . In practice , the expectation under the masking pattern is computed stochastically using a single sample at each epoch .

Summary: In this paper, the authors show that transformer protein language models can learn protein contacts from the unsupervised language modelling objectives. They also show that the residue-residue contacts can be extracted by sparse logistic regression to learn coefficients on the attention heads. One of the advantages of using transformers models is that they do not require an alignment step nor the use of specialized bioinformatics tools (which are computationally expensive). When compared to a method based on multiple sequence alignment, the transformers models can obtain a similar or higher precision.