Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences

We begin by training a series of Transformers on all the sequences in UniParc (33), holding out a random sample of 1 M sequences for validation. We use these models throughout to investigate properties of the representations and the information learned during pretraining.

To comparatively evaluate generalization performance of different language models, we use UniRef50 (34), a clustering of UniParc at 50% sequence identity. For evaluation, a held-out set of 10% of the UniRef50 clusters is randomly sampled. The evaluation dataset consists of the representative sequences of these clusters. All sequences belonging to the held-out clusters are removed from the pretraining datasets.

We explore the effect of the underlying sequence diversity in the pretraining data. Clustering UniParc shows a power law distribution of cluster sizes (35), implying the majority of sequences belong to a small fraction of clusters. Training using a clustering of the sequences results in a reweighting of the masked language modeling loss toward a more diverse set of sequences. We use UniRef (34) to create three pretraining datasets with differing levels of diversity: 1) the low-diversity dataset (UR100) uses the UniRef100 representative sequences, 2) the high-diversity sparse dataset (UR50/S) uses the UniRef50 representative sequences, and 3) the high-diversity dense dataset (UR50/D) samples the UniRef100 sequences evenly across the UniRef50 clusters.

Table 1 presents modeling performance on the held-out UniRef50 sequences across a series of experiments exploring different model classes, number of parameters, and pretraining datasets. Models are compared using the exponentiated cross entropy (ECE) metric, which is the exponential of the model’s loss averaged per token. In the case of the Transformer, this is $2^{{ℒ}_{MLM}}$. ECE describes the mean uncertainty of the model among its set of options for every prediction: ranging from one for an ideal model to 25 (the number of unique amino acid tokens in the data) for a completely random prediction. To measure the difficulty of generalization to the evaluation set, we train a series of n-gram models across a range of context lengths and settings of Laplace smoothing on UR50/S. The best n-gram model has an ECE of 17.18 with context size of four.

As a baseline, we train recurrent long short-term memory (LSTM) bidirectional language models (9), which are state of the art for recurrent models in the text domain. Unlike standard left-to-right autoregressive LSTMs, these models use the whole sequence context, making them comparable to the Transformers we study. We evaluate a small model with ∼25 M parameters and a large model with ∼113 M parameters. Trained on the UR50/S dataset, the small and large LSTM models have an ECE of 14.4 and 13.5, respectively.

We also train two small Transformers, a 12-layer (85.1 M parameters) and six-layer Transformer (42.6 M parameters) on the UR50/S dataset. Both Transformer models have better ECE values (10.45 and 11.79, respectively) than the small and large LSTM models, despite the large LSTM having more parameters. These results show the Transformer enables higher fidelity modeling of protein sequences for a comparable number of parameters.

We train high-capacity 34-layer Transformers (approximately 670 M parameters) across the three datasets of differing diversity. The high-capacity Transformer model trained on the UR50/S dataset outperforms the smaller Transformers, indicating an improvement in language modeling with increasing model capacity. Transformers trained on the two high-diversity datasets, UR50/S and UR50/D, improve generalization over the UR100 low-diversity dataset. The best Transformer trained on the most diverse and dense dataset reaches an ECE of 8.46, meaning that intuitively, the model is choosing among ∼8.46 amino acids for each prediction.

We also train a series of 34-layer Transformer models on 0.1, 1, and 10% of the UR50/S dataset, seeing the expected relationship between increased data and improved generalization performance. Underfitting is observed even for the largest models trained on 100% of UR50/S, suggesting potential for additional improvements with higher capacity models.

Finally, we perform a systematic optimization of model hyperparameters on 100 M parameter models to identify a robust set of hyperparameters. The hyperparameter search is described in detail in SI Appendix, section B. We scale the hyperparameters identified by this search to train a model with ∼650 M parameters (33 layers) on the UR50/S dataset, resulting in the ESM-1b Transformer.

The Transformer neural network represents the identity of each amino acid in its input and output embeddings. The input embeddings project the input amino acid tokens into the first Transformer block. The output embeddings project the final hidden representations back to logarithmic probabilities. The interchangeability of amino acids within a given structural or functional context in a protein depends on their biochemical properties (36). Self-supervision can be expected to capture these patterns to build a representation space that reflects biochemical knowledge.

To investigate if the network has learned to encode physicochemical properties in its representations, we project the weight matrix of the final embedding layer of the network into two dimensions with t-distributed stochastic neighbor embedding (t-SNE) (42). In Fig. 1, the structure of the embedding space reflects biochemical interchangeability with distinct clustering of hydrophobic and polar residues, aromatic amino acids, and organization by molecular weight and charge.

Each protein can be represented as a single vector by averaging across the hidden representation at each position in its sequence. Protein embeddings represent sequences as points in a high dimensional space. Each sequence is represented as a single point, and sequences assigned to similar representations by the network are mapped to nearby points. We investigate how homologous genes are represented in this space.

The structure and function of orthologous genes are likely to be retained despite divergence of their sequences (43). We find in Fig. 2A that training shapes the representation space so that orthologous genes are clustered. Fig. 2A shows a two-dimensional projection of the model’s representation space using t-SNE. Prior to training, the organization of orthologous proteins in the model’s representation space is diffuse. Orthologous genes are clustered in the learned representation space.

We examine whether unsupervised learning encodes biological variations into the structure of the representation space. We apply principal component analysis (PCA) to recover principal directions of variation in the representations, selecting four orthologous genes across four species to look for directions of variation. Fig. 2B indicates that linear dimensionality reduction recovers species and orthology as primary axes of variation in the representation space after training. This form of structure is absent from the representations prior to training.

To quantitatively investigate the structure of the representation space, we assess nearest neighbor recovery under vector similarity queries. If biological properties are encoded along independent directions in the representation space, then proteins corresponding with a unique biological variation are related by linear vector arithmetic. In SI Appendix, Fig. S1, we find that learning improves recovery of target proteins under queries encoded as linear transformations along the species or gene axes.

Remotely homologous proteins have underlying structural similarity despite divergence of their sequences. If structural homology is encoded in the metric structure of the representation space, then the distance between proteins reflects their degree of structural relatedness.

We investigate whether the representation space enables detection of remote homology at the superfamily (proteins that belong to different families but are in the same superfamily) and fold (proteins that belong to different superfamilies but have the same fold) level. We construct a dataset to evaluate remote homology detection using SCOPe (Structural Classification of Proteins—extended) (44). Following standard practices (45), we exclude Rossmann-like folds (c.2 to c.5, c.27 and 28, c.30 and 31) and four- to eight-bladed $\beta$-propellers (b.66 to b.70).

An unsupervised classifier on distance from the query measures the density of homologous proteins in the neighborhood of a query sequence. For each domain, a vector similarity query is performed against all other domains, ranking them by distance to the query domain. For evaluation at the fold level, any domain with the same fold is a positive, any domain with a different fold is a negative, and domains belonging to the same superfamily are excluded. For evaluation at the superfamily level, any domain with the same superfamily is a positive, any domain with a different superfamily is a negative, and domains belonging to the same family are excluded. We report the area under the ROC curve (AUC) for the classifier and Hit-10 (46), which gives the probability of recovering a remote homolog in the 10 highest ranked results.

Table 2 indicates that vector nearest neighbor queries using the representations can detect remote homologs that are distant at the fold level with similar performance to HHblits (15), a state-of-the-art hidden Markov model (HMM) based method. At the superfamily level, where sequence similarity is higher, HMM performance is better, but Transformer embeddings are close. Fast vector nearest neighbor finding methods allow billions of sequences to be searched for similarity to a query protein within milliseconds (47).

An MSA identifies corresponding sites across a family of related sequences (23). These correspondences give a picture of evolutionary variation at different sites within the sequence family. The model receives as input individual sequences and is given no access to the family of related sequences except via learning. We investigate whether the final hidden representations of a sequence encode information about the family it belongs to.

Family information could appear in the network through assignment of similar representations to positions in different sequences that are aligned in the family’s MSA. Using the collection of MSAs of structurally related sequences in Pfam (48), we compare the distribution of cosine similarities of representations between pairs of residues that are aligned in the family’s MSA to a background distribution of cosine similarities between unaligned pairs of residues. A large difference between the aligned and unaligned distributions implies that the representations use shared features for related sites within all the sequences of the family.

Fig. 3A depicts the distribution of cosine similarity values between aligned and unaligned positions within a representative family for the trained model and baselines. Unsupervised learning produces a marked shift between the distributions of aligned and unaligned pairs. Fig. 3 B and C indicate that these trends hold under the constraints that the residue pairs (1) share the same amino acid identity or (2) have different amino acid identities.

We estimate differences between the aligned and unaligned distributions across 128 Pfam families using AUC as a metric of discriminative power between aligned and unaligned pairs. SI Appendix, Table S1 shows a quantitative improvement in average AUC after unsupervised training, supporting the idea that self-supervision encodes information about the MSA of a sequence into its representation of the sequence.

We begin by identifying information about protein structure that is linearly encoded within the representations. The use of linear projections ensures that the information originates in the Transformer representations, enabling a direct inspection of the structural content of representations. By comparing representations of the Transformer before and after pretraining, we can identify the information that emerges as a result of the unsupervised learning.

We perform a fivefold cross validation experiment to study generalization of structural information at the family, superfamily, and fold level. For each of the three levels, we construct a dataset of 15,297 protein structures using the SCOPe database. We partition the structures into five parts, splitting by family, superfamily, and fold accordingly. Fivefold cross validation is performed independently for each of the levels of structural holdout.

To detect information about secondary structure, we fit a logistic regression to the hidden representations using the eight-class secondary structure labels. To detect information about tertiary structure, we fit two separate linear projections to the hidden representations of pairs of positions in the sequence, taking their dot product to regress a binary variable indicating whether the positions are in contact in the protein’s three-dimensional structure. The neural representations are compared to 1) projections of the sequence profile and 2) unsupervised contacts predicted by the CCMpred implementation (51) of direct coupling analysis. MSAs for the baselines are generated from the UniClust30 (52) database using three iterations of search by HHblits. For secondary structure, we report eight-class accuracies. For contact precision, we report top-L long-range precision, that is, the precision of the L (length of the protein) highest ranked predictions for contacts with sequence separation of at least 24 residues.

Table 3 shows results of the cross validation experiments. Prior to pretraining, minimal information about secondary structure and contacts can be detected. After pretraining, projections recover information about secondary structure and long-range contacts that generalizes across families, superfamilies, and folds. Secondary structure prediction eight-class accuracy distributions (SI Appendix, Fig. S2) and long-range contact prediction top-L precision distributions (SI Appendix, Fig. S3) demonstrate that pretraining produces an increase in structural information across the entire distribution of test domains. Table 3 shows that projections of the Transformer representations recover more structure than projections of the sequence profile. For long-range contacts, projections of the best Transformer models have higher precision than contacts predicted by CCMpred across all levels of structural generalization. As the level of structural split becomes more remote, there is little degradation for secondary structure, with performance at the family level similar to the fold level. For long-range contacts, while generalization is reduced at the fold level in comparison to the family level, the best models still capture more structure than the unsupervised baseline. Training with higher diversity sequences (UR50 datasets) improves learning of both secondary structure and long-range contacts, with a more pronounced effect on long-range contacts.

Fig. 4 visualizes three-class secondary structure projections for two domains belonging to held-out folds. Prior to pretraining, projections produce an incoherent prediction of secondary structure. After pretraining, projections recover a coherent prediction with most errors occurring at the boundaries of secondary structure regions. Fig. 5 compares a projected contact map to predictions from CCMpred. Transformer projections recover complex contact patterns, including long-range contacts. Further visualizations of projected contacts for eight randomly selected test proteins are shown in SI Appendix, Fig. S7.

We train deep neural networks to predict secondary structure and contacts from the representations. We choose state-of-the-art neural architectures for both tasks. These downstream models are trained with a supervised loss to predict either the secondary structure or contact map from the pretrained representations. The architecture of the downstream model is kept fixed across experiments with different representations and baselines to enable comparison.

To predict the secondary structure, we replace the linear layer with a deep neural network using the model architecture introduced by the NetSurf method (53). For tertiary structure, we predict the binary contact map from the hidden representation of the sequence. We use a dilated convolutional residual network similar to recent state-of-the-art methods for tertiary structure prediction (54–56).

Table 4 compares the representations for secondary structure prediction. We evaluate models on the CB513 test set (57) and the CASP13 domains (58). For comparison, we also reimplement the NetSurf method. The models are trained on the NetSurf training dataset, which applies a 25% sequence identity holdout with CB513 and a temporal holdout with CASP13. The Transformer features are compared before and after unsupervised pretraining to features from the LSTM baselines. They are also compared to the HMM profiles used by NetSurf. The best Transformer features (71.6%) match the performance of the HMM profiles (71.2%) and exceed the published performance of RaptorX (70.6%) on the same benchmark (53), implying that protein language models can produce features that are directly competitive with sequence profiles for secondary structure prediction.

Table 5 shows the performance of the various representations for predicting top-L long-range contacts across a panel of benchmarks using the RaptorX train set (59). For comparison, we train the same architecture using features from RaptorX (54, 59). The Test (59) and CASP11 (60) test sets evaluate with sequence identity holdout at 25%, the CASP12 (61) test set implements a temporal holdout with the structural training data, and the CASP13 (58) experiment implements a full temporal holdout of both the pretraining and training data. For contact prediction, the best features from representation learning do not achieve comparable performance to the state-of-the-art RaptorX features (50.2 versus 59.4, respectively, on the RaptorX test set).

In the secondary structure benchmarks, Transformer representations produce higher accuracy than the LSTM baselines with comparable numbers of parameters. For contact prediction, Transformer representations yield higher precision than LSTMs, with even the smallest Transformer representations exceeding LSTMs with more parameters. Diversity in the pretraining data also has a strong effect, with the high-diversity datasets providing significant improvements over the low-diversity dataset. Relative performance of the representations is consistent across all four of the contact benchmarks using different holdout methodology.

To investigate the relationship between the language modeling objective and information about structure in the model, linear projections for secondary structure and contacts are fit using the representations from Transformer models taken from checkpoints across their pretraining trajectories. We use the Transformers trained on UR50/S. We fit the projections and evaluate with the train and test split implemented by the first partition of the fold level structural holdout dataset. For each model, Fig. 6 shows a linear relationship between the language modeling objective and information about structure, which is maintained over the course of pretraining. The linear fit is close to ideal for both secondary structure and contacts. A similar experiment is also performed for secondary structure with a deep neural network instead of linear projection by using the NetSurf training sequences and CB513 test set. A linear relationship between secondary structure accuracy and language modeling ECE is also observed for the deep prediction head (SI Appendix, Fig. S4). Thus, for a given model and pretraining dataset, language modeling fidelity measured by ECE is a good proxy for the structural content of the representations. Since performance on the language modeling objective improves with model capacity, this suggests further scale may improve results on structure prediction tasks.

We compare training across evolutionary statistics to training on single protein families. We pretrain separate 12-layer Transformer models on the Pfam multiple sequence alignments of the three most common domains in nature longer than 100 amino acids: the ATP-binding domain of the ATP-binding cassette transporters, the protein kinase domain, and the response regulator receiver domain. We test the ability of models trained on one protein family to generalize secondary structure information within family and out of family by evaluating on sequences with ground truth labels from the family the model was trained on or from the alternate families. The models are evaluated using linear projections. In all cases, the model trained on within-family sequences has higher accuracy than models trained on out-of-family sequences (SI Appendix, Table S2), indicating poor generalization when training on single MSA families. More significantly, the model trained across the full UniParc sequence diversity has a higher accuracy than the single-family model accuracies, even on the same-family evaluation dataset. This suggests that the representations learned from the full dataset are generalizing information about secondary structure learned outside the sequence family.

Current state-of-the-art methods for secondary structure prediction have high accuracies for the eight-class prediction problem. We investigate whether performance can be improved by combining Transformer features with sequence profiles. Table 6 shows that combining the representations with profiles further boosts accuracy, resulting in state-of-the-art performance on secondary structure prediction.

We establish a baseline of performance by reimplementing the Klausen et al. (53) method using the same features, resulting in an accuracy of 71.2% (versus published performance of 72.1%) on the CB513 test set. Then, we add the Transformer features using the direct and avg combination methods; these achieve 0.9 and $2.5\%$ absolute improvement in accuracy respectively. This suggests that the Transformer features contain information not present in the MSA-derived features.

Deep neural networks have enabled recent breakthroughs in the prediction of protein contacts and tertiary structure (54, 56). State-of-the-art neural networks for tertiary structure and contact prediction use deep residual architectures with two-dimensional convolutions over pairwise feature maps to output a contact prediction or distance potential for each pair of residues (54, 56, 59).

A variety of input features, training datasets, and supervision signals are used in state-of-the-art methods. To make a controlled comparison, we fix a standard architecture, training dataset, multiple sequence alignments, and set of base input features for all experiments, to which we add pretrained features from the Transformer model. For the base features, we use the RaptorX feature set, which includes position specific scoring matrix (PSSM), three-state secondary structure prediction, one-hot embedding of sequence, average product correction-corrected Potts model couplings, mutual information, pairwise contact potential, and predicted accessibility. RaptorX was the winning method for contact prediction in the CASP12 and CASP13 competitions (54). The training and evaluation sets are the same as used in the previous section.

Table 7 indicates that addition of Transformer features from the 34-layer model trained on UR50/S consistently produces an improvement across the test sets. The table shows top-L long-range precisions reporting mean and SD over five different model seeds. Direct gives a modest improvement on some test sets. Avg improves over direct, and cov provides further gains. For example, cov produces an absolute improvement of 3.9% on the RaptorX test set and a 1.8% improvement on the CASP13 test set evaluated with temporal holdouts on both fine-tuning and pretraining data. Additional results and metrics for contact prediction are reported in SI Appendix, Table S3.

We analyze the Transformer’s ability to generalize to the fitness landscape of a new protein. Following the protocol introduced in Envision, we use a leave-one-out analysis: to evaluate performance on a given protein, we train on data from the remaining $n-1$ proteins and test on the held-out protein. SI Appendix, Fig. S6 shows that the Transformer’s predictions from raw sequences perform better than Envision on five of the nine tasks.