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: The paper performs a number of analyses centered around the ability of transformer-based language models trained on protein sequence data to learn representations useful for predicting protein secondary and tertiary structure (the latter as contact maps). Specifically, the paper studies several pre-trained transformer models by fitting an L1-penalized logistic regression to amino acid pair contacts. Several experiments are performed to showcase that (i) transformer-based representations can outperform state-of-the art methods based on MSA in terms of contact prediction precision; (ii) that the necessary information for contact predictions in these representations is learned in an unsupervised manner (and not by the logistic regression put on top of these representations); and (iii) that the contact prediction probabilities are reasonably well calibrated.