T RANSFORMER PROTEIN LANGUAGE MODELS ARE UNSUPERVISED STRUCTURE LEARNERS
Anonymous authors
Paper under double-blind review
A BSTRACT
Unsupervised contact prediction is central to uncovering physical, structural, and functional constraints for protein structure determination and design.
For decades, the predominant approach has been to infer evolutionary constraints from a set of related sequences.
In the past year, protein language models have emerged as a potential alternative, but performance has fallen short of state-of-the-art approaches in bioinformatics.
In this paper we demonstrate that Transformer attention maps learn contacts from the unsupervised language modeling objective.
We ﬁnd the highest capacity models that have been trained to date already outperform a stateof-the-art unsupervised contact prediction pipeline, suggesting these pipelines can be replaced with a single forward pass of an end-to-end model.
I NTRODUCTION
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 ﬁt 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 ﬁeld 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 and recurrent LSTM models, the Transformer constructs a pairwise interaction map between all positions in the sequence.
In principle this mechanism has an ideal form to model residue-residue 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 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 identiﬁed 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 signiﬁcant; the advantage holds on average even when Gremlin is given access to an optimized set of multiple sequence alignments incorporating metagenomics data.
We ﬁnd a linear relationship between language modeling perplexity and contact precision.
We also ﬁnd evidence for the value of parameter sharing: the ESM-1b model signiﬁcantly 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.
B ACKGROUND
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 | ∈
for secondary structure prediction in Section 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).
R ELATED W ORK
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 ﬁtting 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.
and reintroduced by Thomas et al.
(2008) and Weigt et al.
Various methods have been developed to ﬁt the underlying Markov Random Field, including meanﬁeld 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 ﬁtting 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 ﬁt with pseudolikelihood maximization to predict pairwise distances with a deep residual network and optimize the ﬁnal 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).
Rives et al.
(2019) were the ﬁrst 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) studied the self-attention mechanism and discovered a correlation between self-attention 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 ﬁrst to ﬁne-tune an LSTM pretrained on protein sequences to ﬁt contacts.
Rao et al.
and Rives et al.
(2020) perform benchmarking of multiple protein language models using a deep residual network ﬁt with supervised learning on top of pretrained language modeling features.
In contrast to previous work on protein language models, we ﬁnd 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.
M ODELS
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. O BJECTIVES
For a set of training sequences, X , Gremlin optimizes the following pseudolikelihood loss, where a single position is masked and predicted from its context.
Since inputs to Gremlin are aligned they all have length L :
[Equation]
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:
[Equation]
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
[Figure]
individual MSAs to massive sequence datasets.
In practice, the expectation under the masking pattern is computed stochastically using a single sample at each epoch.
4. G REMLIN
The log probability optimized by Gremlin is described in section A.3.
Contacts are extracted from the pairwise Gremlin parameters by taking the Frobenius norm along the amino acid dimensions, resulting in an L × L coupling matrix.
Average product correction (APC) is applied to this coupling matrix to determine the ﬁnal predictions (Section A.2).
Gremlin takes an MSA as input.
The quality of the output predictions are highly dependent on the construction of the MSA.
We compare to Gremlin under two conditions.
In the ﬁrst condition, we present Gremlin with all MSAs from the trRosetta training set (Yang et al., 2019).
These MSAs were generated from all of Uniref100 and are also supplemented with metagenomic sequences when the depth from Uniref100 is too low.
The trRosetta MSAs are a key ingredient in the state-of-theart protein folding pipeline.
See Yang et al.
for a discussion on the signiﬁcant impact of metagenomic sequences on the ﬁnal result.
In the second setting, we allow Gremlin access only to the same information as the ESM Transformers by generating MSAs via Jackhmmer on the ESM training set (a subset of Uniref50).
See Section A.5 for Jackhmmer parameters.
4. T RANSFORMERS
We evaluate several pre-trained Transformer models, including ESM-1 (Rives et al., 2019), ProtBertBFD (Elnaggar et al., 2020) and the TAPE Transformer (Rao et al., 2019).
The key differences between these models are the datasets, model sizes, and hyperparameters (major architecture differences described in Table 3).
Liu et al.
(2019) previously showed that these changes can have a signiﬁcant impact on ﬁnal model performance.
In addition to ESM-1, we also evaluate an updated version, ESM-1b, which is the result of a hyperparameter sweep.
The weights of ESM-1b were provided to us by the authors and the differences are described in section A.4.
The Transformer processes inputs through a series of blocks alternating multi-head self-attention and feed-forward layers.
In each head of a self-attention layer, the Transformer views the encoded representation as a
[Table]
set of query-key-value triples.
The output of the head is the result of scaled dot-product attention:
[Equation]
Rather than only computing the attention once, the multi-head approach runs scaled dot-product attention multiple times in parallel and concatenates the output.
Since self-attention explicitly constructs pairwise interactions ( QK T ) between all positions in the sequence, the model can directly represent residue-residue interactions.
In this work, we demonstrate that the QK T pairwise “self attention maps” indeed capture accurate contacts.
4. L OGISTIC R EGRESSION
To extract contacts from a Transformer, we ﬁrst pass the input sequence through the model to obtain the attention maps (one map for each head in each layer).
We then symmetrize and apply APC to each attention map independently.
The resulting maps are passed through an L 1 -regularized logistic regression, which is applied independently at each amino acid pair ( i, j ).
At training time, we only train the weights of the logistic regression; we do not backpropagate through the entire model.
At test time, the entire prediction pipeline can be run in a single forward pass, providing a single end-toend pipeline for protein contact prediction that does not require any retrieval steps from a sequence database.
See Figure 1 for an illustration of this pipeline and Section A.7 for a full description of the logistic regression setup.
R ESULTS
We evaluate models with the 15051 proteins in the trRosetta training dataset (Yang et al., 2019), removing 43 proteins with sequence length greater than 1024 , since ESM-1b was trained with a context size of 1024 .
Of these sequences, Jackhmmer fails on 126 when we attempt to construct
MSAs using the ESM training set (see Section A.5).
This leaves us with 14882 total sequences.
We reserve 20 sequences for training, 20 sequences for validation, and 14842 sequences for testing.
Table 1 shows evaluations of Gremlin, ESM-1, ESM-1b as well as the TAPE and ProtBERT-BFD models.
Conﬁdence intervals are within 0.5 percentage points for all statistics in Tables 1 and 2.
In
Table 1, all Transformer model contact predictors are trained with logistic regression on 20 proteins.
We ﬁnd that with only 20 training proteins ESM-1b has higher precision than Gremlin for short, medium, and long range contacts.
In addition to this set, we also evaluate performance on 15 CASP13 FM Domains in Section A.6.
On average ESM-1b has higher short, medium, and long range precision than Gremlin on all metrics,
[Table]
, s=
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
n=20,
[Table]
=
[Table] [Table]
[Table]
n=20, s=
[Table]
[Table] [Table]
and in particular can signiﬁcantly outperform on MSAs with low effective number of sequences. We also provide a comparison to the bilinear model proposed by Rives et al. The logistic regression model achieves a long-range contact precision at L of 18.6, while the fully supervised bilinear model achieves a long range precision at L of 20.1, an increase of only 1.5 points despite being trained on 700x more structures.
5. A BLATIONS : L IMITING SUPERVISION
While the language modeling objective is fully unsupervised, the logistic regression is trained with a small number of supervised examples.
In this section, we study the dependence of the results on this supervision, providing evidence that the contacts are indeed learned in the unsupervised phase, and the logistic regression is only necessary to extract the contacts.
Top Heads Here we use the logistic regression only to determine the most important heads.
Once they are selected, we discard the weights from the logistic regression and simply average the attention heads corresponding to the top- k weight values.
By taking the single best head from ESM-1b, we come close to Gremlin performance given the same data, and averaging the top-5 heads allows us to outperform Gremlin.
Averaging the top-10 heads outperforms a full logistic regression on all other Transformer models and comes close to Gremlin given optimized MSAs.
Low-N The second variation we consider is to limit the number of supervised examples provided to the logistic regression.
We ﬁnd that with only a single training example , the model achieves a long range top-L precision of 39.2, which is statistically indistinguishable from Gremlin ( p > 0 . 05 ).
Using only 10 training examples, the model outperforms Gremlin on all the metrics.
Since these results depend on the sampled training proteins, we also show a bootstrapped performance distribution using 100 different logistic regression models in Section A.10.
We ﬁnd that withprotein, performance can vary signiﬁcantly, with long range top-L precision mean of 35.6, a median of 38.4, and standard deviation 8.9.
This variation greatly decreases when training on 20 proteins, with a long range top-L precision mean of 40.1, median of 41.1, and standard deviation of 0.3.
See
Figure 11 for the full distribution on all statistics.
MSA Only Finally, we consider supervising the logistic regression only with MSAs instead of real structures.
This is the same training data used by the Gremlin baseline.
To do this, we ﬁrst train
Gremlin on each MSA.
We take the output couplings from Gremlin and mark the top L couplings with sequence separation ≥ 6 in each protein as true contacts, and everything else as false contacts, creating a binary decision problem. When trained on 20 MSAs, we ﬁnd that this model achieves a [Figure] [Figure]
L ong
R a ng e P @ L ESM (6) [Figure]
[Figure]
[Figure]
long range P@L of 39.9, and generally achieves similar long range performance to Gremlin, while still having superior short and medium range contact precision.
5. E NSEMBLING OVER MSA
Transformer models are fundamentally single-sequence models, but we can further boost performance by ensembling predictions from multiple sequences in the alignment.
To do so, we unalign each sequence in the alignment (removing any gaps), pass the resulting sequence through the Transformer and regression, and realign the resulting contact maps to the original aligned indices.
For these experiments, we use the logistic regression weights trained on single-sequence inputs, rather than re-training the logistic regression on multi-sequence inputs.
We also simply take the ﬁrst s sequences in the MSA.
Table 2 shows performance improvements from averaging over 16, 32, and 64 sequences.
To better understand this result, we return to the single-sequence setting and study the change in prediction when switching between sequences in the alignment.
We ﬁnd that contact precision can vary signiﬁcantly depending on the exact sequence input to the model, and that the initial query sequence of the MSA does not necessarily generate the highest contact precision (Figure 9).
Lastly, Alley et al.
(2019) presented a method of ﬁne-tuning where a pretrained language model is further trained on the MSA of the sequence of interest (‘evotuning’).
Previously this has only been investigated for function prediction and for relatively low-capacity models.
We ﬁne-tune the full
ESM-1b model (which has 50x more parameters than UniRep) on 380 protein sequence families.
We ﬁnd that after 30 epochs of ﬁne-tuning, long range P@L increases only slightly, with an average of 1.6 percentage points (Figure 15).
5. P ERFORMANCE D ISTRIBUTION
Although our model is, on average, better than Gremlin at detecting contacts, the performance distribution over all sequences in the dataset is still mixed.
ESM-1b is consistently better at extracting short and medium range contacts (Figure 7), but only slightly outperforms Gremlin on long range contacts when Gremlin has access to Uniref100 and metagenomic sequences.
Figure 2 shows the distribution of long range P@L for ESM-1b vs. Gremlin.
Overall, ESM-1b has higher long range
P@L on 55% of sequences in the test set.
In addition, we examine the relationship between MSA depth and precision for short, medium, and long range contacts (Figure 3).
Although our contact prediction pipeline does not make explicit use of MSAs, there is still some correlation between MSA depth and performance, since MSA depth is a measure of how many related sequences are present in the ESM-1b training set.
We again see that ESM-1b consistently outperforms Gremlin at all MSA depths for short and medium range sequences.
We also conﬁrm that ESM-1b outperforms Gremlin for long range contact extraction for sequences with small MSAs (depth < 1000 ).
ESM-1b also outperforms Gremlin on sequences with
7 .10 .12 .log
2 (depth) 0 .0 .
P @
L 6 ≤ s <7 .10 .12 .log
2 (depth) 0 .0 .
.12 ≤ s <7 .10 .12 .log
2 (depth) 0 .
.24 ≤ s
ESM-1b
Gremlin the very largest MSAs (depth > 16000 ), which is consistent with prior work showing that Gremlin performance plateaus for very large MSAs and suggests that ESM-1b does not suffer from the same issues (Anishchenko et al., 2017).
5. L OGISTIC R EGRESSION W EIGHTS
In Section 5.1 we show that selecting only a sparse subset of the attention heads can yield good results for contact prediction.
Overall, the L 1 -regularized logistic regression identiﬁes 102 /heads as being predictive of contacts (Figure 6b). Additionally, we train separate logistic regressions to identify contacts at different ranges (local, short, medium, and long range). These regressions identify an overlapping, but non-identical set of useful attention heads. Two attention heads have the top-10 highest weights for detecting contacts at all ranges.
One attention head is highly positively correlated with local contacts, but highly negatively correlated with long range contacts.
Lastly, we identify a total of 104 attention heads that are correlated (positively or negatively) with contacts at only one of the four ranges, suggesting that particular attention heads specialize in detecting certain types of contacts.
5. P ERPLEXITY VS . C ONTACT P RECISION
Figure 2 explores the relationship between performance on the masked language modeling task (validataion perplexity) and contact prediction (Long Range P@L).
A linear relationship exists between validation perplexity and contact precision for each model.
Furthermore, for the same perplexity, the 12-layer ESM-1 model achieves the same long range P@L as the 34 layer ESM-1 model, suggesting that perplexity is a good proxy task for contact prediction.
ESM-1 and ESM-1b models are trained with different masking patterns, so their perplexities cannot be directly compared, although a linear relationship is clearly visible in both.
ESM-
and ESM-1b have a similar number of parameters; the key difference is in their hyperparameters and architecture.
The models shown have converged in pre-training, with minimal decrease in perplexity (or increase in contact precision) in the later epochs.
This provides clear evidence that both model scale and hyperparameters play a signiﬁcant role in a model’s ability to learn contacts.
5. C ALIBRATION , F ALSE P OSITIVES , AND R OBUSTNESS
One concern with large neural networks is that, while they may be accurate on average, they can also produce spurious results with high conﬁdence.
We investigate this possibility from several perspectives.
First, we ﬁnd that logistic regression probabilities are close to true contact probability (mean-squared error = 0 . 014 ) and can be used directly as a measure of the model’s conﬁdence (Figure 12a).
Second, we analyze the false positives that the model does predict.
We ﬁnd that these are very likely to be within a Manhattan distance of 1-4 of a true contact (Figure 12b).
This suggests that false positives may arise due to the way a contact is deﬁned (Cb-Cb distance within 8 angstroms), and could be marked as true contacts under a different deﬁnition (Zheng & Grigoryan, 2017).
Further, when we explore an example where the model’s predictions are not near a true contact, we see that {+[Figure] Head   
 L+} a y e r
{+Local 0 3 6 9 12 
+}
Head
{+Short 0 3 6 9 12 
+}
Head
{+Medium 0 3 6 9 12 
+}
Head
Long −−   the example in question is a homodimer, and that the model is picking up on inter-chain interactions (Figure 13a).
While these do not determine the structure of the monomer, they are important for its function (Anishchenko et al., 2017).
Third, we test the robustness of the model to insertions by inserting consecutive alanines at the beginning, middle, or end of 1000 randomly chosen sequences.
We ﬁnd that ESM-1b can tolerate up to 256 insertions at the beginning or end of the sequence and up to 64 insertions in the middle of the sequence before performance starts to signiﬁcantly degrade.
This suggests that ESM-1b learns a robust implicit alignment of the protein sequence.
See Section A.12 for more details.
5. MSA G ENERATION
Wang & Cho (2019) note that Transformers trained with the MLM objective can be used generatively.
Here, we consider whether generations from ESM-1b preserve contact information for a given protein.
The ability to generate sequences that preserve this information is a necessary condition for generation of biologically active proteins.
We perform this evaluation by taking an input protein, masking out several positions, and re-predicting them.
This process is repeated 10000 times to generate a pseudo-MSA for the input sequence (see Algorithm 1).
We feed the resulting MSA into Gremlin to predict contacts.
Over all sequences from our test set, this procedure results in a long range contact P@L of 14.5.
Figure 16 shows one example where the procedure works well, with Gremlin on the pseudo-MSA having long range P@L of 52.2.
For comparison, the standard
ESM-1b pipeline achieves a contact precision of 76.7 on this example.
D ISCUSSION
Transformer protein language models trained with an unsupervised objective learn the tertiary structure of a protein sequence in their attention maps.
Residue-residue contacts can be extracted from the attention by sparse logistic regression.
Attention heads are found that specialize in different types of contacts.
An ablation analysis conﬁrms that the contacts are learned in the unsupervised phase, and that the logistic regression is only necessary to extract the part of the model that represents the contacts.
These results have implications for protein structure determination and design.
The initial studies proposing Transformers for protein language modeling showed that representation learning could be used to derive state-of-the-art features across a variety of tasks, but were not able to show a beneﬁt in the fully end-to-end setting (Rives et al., 2019;
Rao et al., 2019; Elnaggar et al., 2020).
For the ﬁrst time, we show that protein language models can outperform state-of-the-art unsupervised structure learning methods that have been intensively researched and optimized over decades.
Finally, we establish a link between language modeling perplexity and unsupervised structure learning.
A similar scaling law has been observed previously for supervised secondary structure prediction (Rives et al., 2019), and parallels observations in the NLP community (Kaplan et al., 2020;
Brown et al., 2020).
Evidence of scaling laws for protein language modeling support future promise as models and data continue to grow.
A A PPENDIX A. N OTATION
In the ﬁgures, we report contact precision in the range of 0 .
to 1 . 0 .
In the text and in the tables, we report contact precision in terms of percentages, in the range of 0 to 100 .
A. A VERAGE P RODUCT C ORRECTION (APC)
In protein contact prediction, APC is commonly used to correct for background effects of entropy and phylogeny (Dunn et al., 2008).
Given an L × L coupling matrix F , APC is deﬁned as
[Equation]
Where F i , F j , and F are the sum over the i -th row, j -th column, and the full matrix respectively.
We apply APC independently to the symmetrized attention maps of each head in the Transformer.
These corrected attention maps are passed in as input to a logistic regression.
A. G REMLIN I MPLEMENTATION D ETAILS
Gremlin is trained by optimizing the pseudolikelihood of W and V , which correspond to pairwise and individual amino acid propensities.
The pseudolikelihood approximation models the conditional distributions of the original joint distribution and can be written:
[Equation]
subject to the constraint that W ii = 0 for all i , and that W ijab is symmetric in both sequence ( i, j ) and amino acid ( a, b ).
Additionally, Gremlin uses a regularization parameter that is adjusted based on the depth of the MSA.
A. ESM-1 I MPLEMENTATION D ETAILS
The original ESM-1 models were described in (Rives et al., 2019).
ESM-1 is trained on Uniref in contrast to the TAPE model, which is trained on Pfam (Finn et al., 2014) and the ProtBERT-BFD model, which is trained on Uniref100 and BFD100 (Steinegger et al., 2019).
Recently, the authors have made ESM-1b available, which is a robustly optimized version of the largest ESM-1 model.
ESM-1b is the result of an extensive hyperparameter sweep that was performed on smaller 12 layer models.
ESM-1b is the result of scaling up that model to 33 layers.
Compared to ESM-1, the main changes in ESM-1b are: higher learning rate; dropout after word embedding; learned positional embeddings; ﬁnal layer norm before the output; and tied input/output word embeddings.
A. J ACKHMMER D ETAILS
We use Jackhmmer version 3.3.1 with a bitscore threshold of 27 and 8 iterations to construct MSAs from the ESM training set.
The failures on 126 sequences noted in Section 4.4 result from a segmentation fault in hmmbuild after several iterations (the number of successful iterations before the [Table] bilinear) [Table]
segmentation fault varies depending on the input sequence).
Since we see this failure for less than 1% of the dataset we choose to ignore these sequences during evaluation.
Additionally, we evaluated alternate MSAs by running Jackhmmer until a Neff of 128 was achieved (with a maximum of 8 iterations), a procedure described by Zhang et al.
This resulted in very similar, but slightly worse results (average long range P@L 29.3, versus 31.3 when always using the output of the eighth iteration).
We therefore chose to report results using the 8 iteration maximum.
A. R ESULTS ON CASP
In Table 4 we report results on the 15 CASP13 Free Modeling targets for which PDBs were publicly released.
The speciﬁc domains evaluated are: T0950-D1, T0957s2-D1, T0960-D2, T0963-D2,
T0968s1-D1, T0968s2-D1, T0969-D1, T0980s1-D1, T0986s2-D1, T0990-D1, T0990-D3, T1000D2, T1021s3-D1, T1021s3-D2, T1022s1-D1.
ESM-1b is able to outperform Gremlin, and simply averaging the top-10 heads of ESM-1b has comparable performance to Gremlin.
In addition, we compare our logistic regression model to the bilinear contact prediction model proposed by Rives et al.
This model trains two separate linear projections of the ﬁnal representation layer and computes contact probabilities via the outer product of the two projections plus a bias term, which generates the following unnormalized log probability: [Equation] Here x is a sequence-length vector of features in R L × d . Each W i is a matrix in R d × k , where k is a hyperparameter controlling the projection size.
We train this model in both the limited supervision ( n = 20 ) and full supervision ( n = 14257 ) setting.
For the limited supervision setting, we use the same 20 proteins used to train the sparse logistic regression model. For the full supervision setting we generate a 95/5% random training/validation split of the 15008 trRosetta proteins with sequence length ≤ 1024 . We performed independent grid searches over learning rate, weight decay, and hidden size for the two settings. For the n = 20 setting, we found a learning rate of 0.001, weight decay of 10.0, and projection size of 512 had best performance on the validation set.
{+For the n = 14257 setting we found a learning rate of 0.001, weight decay of 0.01, and projection size of 512 had best performance Gremlin 0 .0 .0 .
 E S M - 1b Long Range P@L 0 .0 .0 .0 .0 .1 .+}
Gremlin
Long Range P@L/ neff
[Figure]
[Figure]
[Figure]
{+Gridsearch 0 3 6 
+}
Head
L a y e r −−−  ( [Figure] ) Weights on the validation set. All models were trained to convergence to maximize validation long range
P@L with a patience of 10.
The n = 20 models were trained with a batch size of 20 (i.e. 1 batch = 1 epoch) and the n = 14257 models were trained with a batch size of 128.
The bilinear model performs very poorly in the limited supervision setting, worse than simply taking the top-1 attention head.
With full supervision, it moderately outperforms the logistic regression for an increase in long range P@L of 1.5 while using 700x more data.
In Figure 5 we display results on the 15 FM targets colored by effective number of sequences.
ESM-1b shows higher precision at L and L/5 on average, and is sometimes signiﬁcantly higher for sequences with low Neff.
Since ESM-1b training data was generated prior to CASP13, this suggests
ESM-1b is able to generalize well to new sequences.
A. L OGISTIC R EGRESSION D ETAILS
Given a model with M layers, H heads, and an input sequence x of length L , let A mh be the L × L contact map from the h -th head in the m -th layer.
We ﬁrst symmetrize this map and apply APC and let a mhij be the coupling weight between sequence position i and j in the resulting map.
Then we deﬁne the probability of a contact between positions i and j according to a logistic regression with parameters β :
[Equation]
To ﬁt β , let D be a set of training proteins, k be a minimum sequence separation, and λ be a regularization weight.
The objective can then be deﬁned as follows:
[Equation]
[-Gridsearch 0 3 6 
-]
We ﬁt the parameters β via scikit-learn (Pedregosa et al., 2011) and do not backpropagate the gradients through the attention weights.
In total, our model learns MH + 1 parameters, many of which are zero thanks to the L 1 regularization.
There are three hyperparameters in our training setup: the number of proteins in our training set D , the regularization parameter λ , and the minimum sequence separation of training contacts k .
We ﬁnd that performance improves signiﬁcantly when increasing the D from 1 protein to 10 proteins, but that the performance gains drop off when D increases from 10 to 20 (Figure 1).
Through a hyperparameter sweep, we determined that the optimal λ is 0.15.
We ﬁnd that ignoring local contacts ( | i − j | < 6 ) is also helpful.
Therefore, unless otherwise speciﬁed, all logistic regressions are trained with |D| = 20 , λ = 0 .
15 , k = 6 .
See Figure 6a for a gridsearch over the number of training proteins and regression penalty.
We used 20 training proteins and 20 validation proteins for this gridsearch.
Figure 6b shows the weights of the ﬁnal logistic regression used for ESM-1b.
A. P ERFORMANCE D ISTRIBUTION
Figure 7 shows the full distribution of performance of ESM-1b compared with Gremlin.
When we provide Gremlin access to Uniref100, along with metagenomic sequences, ESM-1b still consistenly outperforms Gremlin when extracting short and medium range contacts.
For long range contacts,
Gremlin is much more comparable, and has higher contact precision on 47% of sequences.
With access to the same set of sequences, ESM-1b consistently outperforms Gremlin in detecting short, medium, and long range contacts.
This suggests that ESM-1b can much better extract information from the same set of sequences and suggests that further scaling of training data may improve ESM1b even further.
This analysis is further borne out in Figure 8.
Given the same set of sequences, ESM-1b outperforms
Gremlin on average for short, medium, and long-range contacts regardless of the depth of the MSA generated from the ESM-1b training set.
Additionally, we ﬁnd that ESM-1b can provide varying contact maps for different sequences in the
MSA (Figure 9).
This is not possible for Gremlin, which is a family-level model.
We leverage this in a fairly simple way to provide a modest boost to the contact precision of ESM-1b (Section 5.2).
A. S ECONDARY S TRUCTURE
In Section 5.4 we show that some heads that detect local contacts (which often correspond to secondary structure) are actually negatively correlated with long range contacts.
We test ESM-1b’s ability to detect secondary structure via attention by training a separate logistic regression on the [Figure]
Netsurf dataset (Klausen et al., 2019).
As with the logistic regression on contacts, we compute attentions and perform APC + symmetrization.
To predict the secondary structure of amino acid i , we feed as input the couplings a mhij for each layer m , for each head h , and for j ∈ [ i − 5 , i + 5] , for a total of 7260 input features.
Using just 100 of the 8678 training proteins, we achieve 79.3% accuracy on 3-class secondary structure prediction on the CB513 test set (Cuff & Barton, 1999).
Figure 10 shows the importance of each layer to predicting the three secondary structure classes.
There are
[Figure]
ong
R a
[Figure]
e P @
L
[Figure]
R a
[Figure]
e [Figure]
L
[Figure]
P @
L
[Figure]
spikes in different layers for all three classes, indicating that particular heads within those layers are specializing in detecting speciﬁc classes of secondary structure.
Figure 10 shows importance of each
Transformer layer to predicting each of the three secondary structure classes.
We see that, as with contact prediction, the most important layers are in the middle layers (14-20) and the ﬁnal layers (29-33).
Some layers spike more heavily on particular contact classes (e.g. layer 33 is important for all classes, but particularly important for β -strand prediction).
This suggests that particular heads within these layers activate speciﬁcally for certain types of secondary structure.
A. B OOTSTRAPPED L OW -N C ONFIDENCE I NTERVAL
Section 5.1 shows results from Low-N supervision on 1, 10, and 20 proteins.
Since performance in this case depends on the particular proteins sampled we use bootstrapping to determine a conﬁdence interval for each of these estimates.
Using the full training, validation, and test set of 14882 proteins, we train 100 logistic regression models using a random sample of N proteins, for N = 1, 10, and 20.
Each model is then evaluated on the remaining 14882 − N proteins.
The full distribution of samples can be seen in Figure 11.
The conﬁdence interval estimates for long range precision at L with 1, 10, and 20 training proteins are: 35 .
6 ± 1 . 8 , 40 .
6 ± 0 . 1 , and 41 . 0 ± 0 .
1 respectively.
[Figure]
[Figure]
[Figure]
[Figure]
A. M ODEL C ALIBRATION AND F ALSE P OSITIVES
Vig et al.
(2020) suggested that the attention probability from the TAPE Transformer was a wellcalibrated estimator for the probability of a contact.
In Figure 12a we examine the same with the logistic regression trained on the ESM-1 and ESM-1b models.
We note that ESM-1b, in addition to being more accurate overall than Gremlin, also provides actual probabilities.
We ﬁnd that as with model accuracy, model calibration increases with larger scale and better hyperparameters.
The 6, 12, and 34 layer ESM-1 models have mean-squared error of 0.074, 0.028, and 0.020 between predicted and actual contact probabilities, respectively. ESM-1b has a mean squared error of 0.014.
Mean squared error is computed between contact probabilites split into 20 bins according to the scikit-learn calibration curve function.
It is therefore reasonable to use the logistic regression probability as a measure of the model’s conﬁdence.
In the case of false positive contacts we attempt to measure the Manhattan distance between the coordinates of predicted contacts and the nearest true contact (Figure 12b).
We observe that the
Manhattan distance between the coordinates of false positive contacts are often very close (Manhattan distance between 1-4) to real contacts, and that very few false positives have a Manhattan distance ≥ 10 from a true contact.
With a threshold contact probability of 0.5, 83.8% of proteins have at least one predict contact with Manhattan distance > 4 to the nearest contact.
This drops to 71.7% with a threshold probability of 0.7, and to 52.5% with a threshold probability of 0.9.
Figure 13 highlights two modes for ESM-1b where signﬁcant numbers of spurious contacts are predicted.
Figure 13a shows one example where the model does appear to hallucinate contacts around residues 215 and 415, which do not appear in the contact map for this protein.
However, this protein is a homodimer and these contacts are present in the inter-chain contact map.
This suggests that some ‘highly incorrect’ false positives may instead be picking up on inter-chain contacts. Figure 13b shows an example of a repeat protein, for which evolutionary coupling methods are known to pick
Insertion Length 0 .0 .0 .0 .0 .0 .0 .
L ong
R a ng e P @ L
Average model
TAPE
ESM-1b location beginning midpoint end    8 {+16 32 64 
+}
Insertion Length 0 .0 .0 .0 .0 {+.0 .
+}
.1a up on additional ‘bands’ of contacts (Espada et al., 2015; Anishchenko et al., 2017).
Multiple bands are visible in the Gremlin contact map, while only the ﬁrst band, closest to the diagonal, is visible in the ESM-1b contact map. More analysis would be necessary to determine the frequency of these modes, along with additional potential modes.
A. A LIGNMENT
One hypothesis as to the beneﬁt of large language models as opposed to simpler Potts models is that they may be able to learn an implicit alignment due to their learned positional embedding.
For a Potts
Model, an alignment enables a model to relate positions in the sequence given evolutionary context despite the presence of insertions or deletions.
We test the robustness of the model to insertions by inserting consecutive alanines at the beginning, middle, or end of 1000 randomly chosen sequences with initial sequence length <
(we limit initial sequence length in order to avoid out-of-memory issues after insertion).
We ﬁnd that ESM-1b can tolerate up to 256 insertions at the beginning or end of the sequence and up to 64 insertions in the middle of the sequence before performance starts to signiﬁcantly degrade.
This suggests that ESM-1b learns a robust implicit alignment of the protein sequence.
On the other hand, we ﬁnd that the TAPE Transformer is less robust to insertions.
On one sequence (pdbid: 1a27), we ﬁnd the TAPE Transformer drops in precision by 12 percentage points after adding just 8 alanines to the beginning of the sequence, while ESM-1b sees minimal degradation untilalanines are inserted.
We hypothesize that, because TAPE was trained on protein domains, it did not learn to deal with mis-alignments in the input sequence.
A. E VOLUTIONARY F INETUNING D ETAILS
We ﬁnetuned each model using a learning rate of 1e-4, 16k warmup updates, an inverse square root learning rate schedule, and a maximum of 30 epochs.
[-ESM-1b location beginning midpoint end    8 16 32 64 
-]
This resulted in a varying number of total updates depending on the size of the MSA, with larger MSAs being allowed to train for more updates.
This should ideally help prevent the model from overﬁtting too quickly on very small
MSAs.
We use a variable batch size based on the length of the input proteins, ﬁxing a maximum of 16384 tokens per batch (so for a length 300 protein this would correspond to a batch size of 54).
We use MSAs from trRosetta for ﬁnetuning all proteins with the exception of avGFP, where we use the same set of sequences from Alley et al. [Figure] Evotuning [Figure]
Evotuned
[Figure]
No Evotuning
Evotuned
[Figure]
A. MSA G ENERATION
Result:
Generated MSA input // protein sequence curr = input // optionally, the input can be repeated for batching for 0 ≤ i < 10000 do masked = mask 20% of positions in curr; pred = model(masked); curr[masked positions] = pred masked positions];
MSA.append(curr); if random () < 0 .
1 then curr = input; end end
Algorithm 1: Quickly generate a pseudo-MSA from an input sequence.
Algorithm 1 presents the algorithm used to generate pseudo-MSAs from ESM-1b.
Each pseudoMSA is passed to GREMLIN in order to evaluate the preservation of contact information (Figure 16).