ESM-2 (evolutionary-scale prediction of atomic level protein structure with a language model)

By Ramith Hettiarachchi ⋅ Published 2023-02-10

Highlights of the ESM-2 Paper

1. Introduction

1.1 — Structure and function are hidden in sequences

Biological properties of proteins influence which positions in a sequence can undergo mutations. From these observations, we can identify evolutionary patterns such as coevolution and conservation of amino acids. These patterns can help us infer properties of protein function and structure.⁴⁴ In essence, information about protein structure and function is hidden in sequences.

Usually, we align sequences before drawing conclusions about function and structure. This intermediate representation, known as a multiple sequence alignment (MSA), has a high time complexity because we must first search for related sequences and then align them.⁵⁵ The search process in AlphaFold and RoseTTAFold pipelines can take more than ten minutes.

What if we can get rid of this intermediate representation? That’s one aspect this paper accomplishes.

1.2 — Large language models (LLMs)

Historically, language models were pretrained using objectives such as predicting the next word in a sentence. Devlin et al.‘s BERT showed that masking some input tokens and predicting them—the masked-language-model objective, or MLM—is an effective pretraining strategy.⁶⁶ Unlike left-to-right pretraining, the MLM objective lets the representation combine both left and right context, enabling a deeply bidirectional Transformer (paraphrased from the BERT paper).

1.3 — Contributions

Inspired by this widely adopted strategy, the authors hypothesise that filling in missing amino acids might produce representations rich enough to infer structure. They therefore scale protein language models from 8 million parameters up to 15 billion parameters. Doing so reveals the following:

Atomic-level structure prediction directly from sequence.
A strong correlation between perplexity and structure-prediction accuracy.
Up to 60× faster inference.
No need to search for related sequences.

Because this approach improves speed by one to two orders of magnitude and does not need an MSA, the authors expand structure prediction to the much larger and more diverse space of metagenomic proteins. In summary, they:

Predict structures for more than 617 million sequences in MGnify90.⁷⁷ This took two weeks on 2,000 GPUs. 🤯
- Of those structures, 225 million are high-confidence predictions.
  - 76.8% are disjoint from UniRef90 at 90% sequence identity.
  - 12.6% have no experimental ground truth.

Open ESM Metagenomic Atlas explorer in a new tab.

2. Method

2.1 — How does structure emerge from a language model trained on sequences?

The ESM-2 language model is trained with approximately 65 million unique sequences.⁸⁸ The sequences were sampled with even weighting across approximately 43 million UniRef50 training clusters. Because the MLM objective asks the model to predict missing amino acids using their neighboring context, the model must learn interdependencies between amino acids. Previous work—[1] and [2]—showed that Transformer models trained with MLM on protein sequences develop attention patterns that correspond to residue-residue contact maps.⁹⁹ I plan to write a separate article about this work.

After training the language model, the authors use the approach from [2] to compute contact maps from attention patterns. A logistic regression identifies contacts as follows.

Annotated contact-prediction pipeline from Transformer attention maps — Attention maps from a masked protein language model are symmetrized and passed to logistic regression to predict residue-residue contacts. Source figure: Rao et al. (2020) annotations added for this article.

2.2 — What about atomic-level structure? (Enter ESMFold)

The authors extract contact maps from attention patterns, but predicting atomic spatial coordinates requires an equivariant Transformer. They use the structure moduleAlphaFold structure module and invariant point attention. introduced in AlphaFold. This module projects atomic spatial coordinates from the language model’s internal representation. The complete architecture is called ESMFold.

Steps in ESMFold

Process the sequence through ESM-2.
Pass the representation learned by ESM-2 through a series of folding blocks. Each block sequentially updates a sequence representation and a pairwise representation.
Pass the result to the structure module.
Repeat with three recycling steps. View annotated code.

ESMFold architecture with ESM-2, folding trunk, structure module, and recycling — ESMFold architecture: ESM-2 representations pass through a folding trunk and an AlphaFold-derived structure module, with recycling. Source figure: Lin et al. (2022) annotations added for this article.

The supporting figures connect the structure-module diagram to implementation details. Select any image to open the full-resolution version.

Annotated AlphaFold structure module and invariant point attention diagram — AlphaFold’s structure module and invariant point attention. Source figures: Jumper et al. (2021).

Annotated ESMFold folding-trunk and triangular-attention source code — Mapping the folding trunk to its triangular self-attention block.

Annotated ESMFold trunk code showing folding, structure prediction, and recycling — Folding-block iteration, the structure-module call, and recycling.

Training: To train the structure model to obtain spatial coordinates, the authors use experimentally determined structures from the Protein Data Bank—approximately 25,000 clusters covering around 325,000 structures. This is augmented with 12 million structures predicted by AlphaFold2.¹⁰¹⁰ During training, predicted structures are sampled 75% of the time and real structures 25% of the time.

Evaluation: 194 CAMEO proteins and 51 CASP14 proteins.

This language model based approach vastly simplifies the usual SOTA structure prediction process by eliminating the need for the following ,

External evolutionary databases
Multiple sequence alignments (MSAs)
Templates

For example, AlphaFold requires access to these resources.

3. Results

3.1 — How well does it predict structures?

As mentioned before, they evaluate performance on CAMEO and CASP14 proteins and check how well the structure was predicted using the TM-Score.

In predicting the structure just by single sequences, ESMFold achieves very good performance compared to AlphaFold and RoseTTAFold.

ESMFold, AlphaFold2, and RoseTTAFold TM-score comparisons on CAMEO and CASP14 — Single-sequence and full-pipeline TM-score comparisons on CAMEO and CASP14; point color encodes perplexity. Source: Figure 2B in Lin et al. (2022) annotations added for this article.

3.2 — How important is the language model in the pipeline?

The key question that arises is how important is the representation learnt by the LM for the task of structure prediction . To quantify this we need several metrics.

First, we need to characterize how well the language model understands protein sequences. This is where perplexity comes in. We already have the TM-score to measure how closely a predicted structure matches the ground truth.

Thus, the graph to the right in Fig. 2B shows that,

High ESMFold TM-Scores have low perplexity scores (numerically speaking, on CAMEO, Pearson correlation coefficient is −0.55 and in CASP14 it’s −0.67)

How can we achieve lower perplexity?

We now know that a better language model representation—with lower perplexity—leads to better structure prediction. How can we obtain a better representation? 🤔 Is scaling all you need?

To answer this question, authors explore the effect of scaling and look at what happens to the following :

Precision @ L
Change in perplexity

Scatter plots comparing long-range contact precision across ESM-2 model sizes — Effect of scaling ESM-2 on long-range contact precision; point color shows the change in perplexity. Source: Figure 1D in Lin et al. (2022) annotations added for this article.

The authors plot how long-range precision at L changes when moving from a smaller model on the x-axis to a larger model on the y-axis. Points above the diagonal suggest that scaling improves long-range precision at L for some proteins.

Is scaling the answer?

It is not that simple. Although P@L increases with scale for some proteins, the number of evolutionarily related sequences tells another story. Language models struggle when less relevant training data is available for a query. This is intuitive—more study, better results—but does it reflect memorization rather than understanding?

Long-range contact precision plotted against the number of related sequences — Long-range precision at L as a function of the number of related sequences for several ESM-2 model sizes. Source: Lin et al. (2022) annotations added for this article.

The authors could have used a different color scheme: white points indicate no change in perplexity, while a point on the diagonal indicates no improvement in long-range precision at L.¹¹¹¹ It would be useful if the plot made the proportion of proteins with improved performance directly visible, although that might make it too cluttered.

3.3 — What about prediction speed?

A protein with 384 residues on one NVIDIA V100 GPU: 14.2 seconds.¹²¹² This is a 6× speedup compared with AlphaFold2.
Shorter sequences 60x speedup

Log-scale inference time comparison for ESMFold, AlphaFold, and RoseTTAFold — Inference time versus sequence length. ESMFold is fastest for shorter sequences, while its pair representation makes long sequences more expensive. MSA search time for the other methods is not included.

3.4 — Comparison with other protein language models

Table comparing ESM-2 sizes and other protein language models — Validation perplexity, long-range contact precision, and structure-prediction results across ESM-2 sizes and comparison protein language models. Source: Lin et al. (2022) annotations added for this article.

4. Conclusion

It’s remarkable that these authors scale protein language models and it has resulted in learning structure hidden through databases of sequences, and thus we do not need to depend onto the MSA.

Is it because, the model has learnt to obtain the signal which we previously obtained through MSAs? What can we tell about the performance of sequences that had less number of evolutionary sequences in training data? why does it still struggle to obtain decent performance. It would be very interesting to analyze these directions.

Thanks for reading this, hope you found it useful. If you have any suggestions/ comments please share below.

References

KUDOSDon’t
MoveThanks!