Protein Structure Prediction Using Large Language Models

A quick hitter post on ESMFold, an LLM trained on tens of millions of protein sequences.

The model architecture of ESMFold

ESMFold is a model that adapts to proteins the same technology underpinning ChatGPT and Claude to simplify and speed up the process of protein folding prediction.

The Role of MSA’s and Evolutionary Information in Protein Structure Prediction

Protein structure prediction using machine learning is an extremely powerful technique that maps protein amino acid residue sequences to the 3D folded structure of the protein. Traditionally, these protein folding models need a little bit more context than just the protein sequence. They also need information about evolutionarily related proteins (homologs) to the protein in question. The main way evolutionary information is handed to protein folding models is through what are called multi-sequence alignments. Namely, the homolog structures (templates) and sequences (multi-sequence alignment). Multi-sequence alignments (MSA) in looking at many different protein sequences with common ancestors, allows us to track how protein sequences have mutated as they’ve evolved. Which residues change and which residues change together (co-varied) can give the model huge hints into which residues are in contact in the final folded structure. While the MSA and the model’s processing of the MSA is crucial for accurate structure prediction, it is an expensive process. Not only do massive protein sequence databases have to be searched and their sequences aligned, but this must be done during inference for each sequence we want a structure for. This database searching phase can take up to 10 minutes for models, like AlphaFold2. Moreover, models like AlphaFold2 must be more memory and compute intensive to process MSA’s.

Where Large Language Models Come In

Since hints to protein structure and function are embedded in their amino acid sequences and these sequences can be represented as text, why can't we use the best tool we know for uncovering and understanding patterns in text: language models?

A key insight from natural language understanding is that a word's meaning can be determined from the surrounding context in which it appears. Large language models take advantage of this insight by being trained using self-supervised learning. In self-supervised learning, a language model can be trained just on text (with no labels) by either predicting the next word given the previous ones or by removing random words then trying to predict these "missing" words using surrounding words as context. This "fill-in-the-blank" task forces the language model to learn a representation that captures a lot of the patterns and meanings of text and can be used as jumping off point for natural language tasks like sentiment analysis of a paragraph or question answering. Often when training language models using SSL, the larger the training set the better the learned representation. That is why training models like GPT3 on the entire text of the internet has given us impressive results.

Just like the internet gave us a massive amount of text to train LLM's on metagenomic sampling and protein databases has given us a huge number of protein sequences ripe for processing with LLM's. With metagenomic sampling, scientists can now take samples from soil or the ocean extract the dna from the organisms they find and then work forwards to figure out what protein sequences the genes in the DNA express. We now have massive databases of hundreds and hundreds of millions of proteins.

Since LLM's take all the text on the internet and predict the next word, why don't we take all the protein sequences we know and instead of filling in the missing word, let's fill in the missing amino acid? This is essentially the question that ESMFold asks. If this LLM's for protein sequences idea works, we can eliminate the need for MSA's, which results in much faster model: 60x speedup in inference forward pass and also eliminating up to 10 minutes needed to generate MSA for each protein input.

The ESMFold Model

The ESMFold architecture is split into two parts: the ESM-2 language model, which is trained to predict “missing” amino acids on millions of protein sequences, where we randomly drop some of their amino acids; next the representations are processed by the “folding head” module, architected similarly to AlphaFold, which after further refining the language model representation, predicts a structure for the protein. Let’s dive into both parts.

ESM-2 Language Model

The ESM-2 model is a transformer protein language model, which is essentially the same underlying technology that ChatGPT, Claude, and Gemini use, but adapted to be used for proteins instead of english and other natural language. To understand ESM-2, let’s do a quick dive into two language models it is based off of: transformers and BERT

The original transformer model

The first is the original transformer model, the model that started it all. Still, the key foundational model that inspired the models underlying ChatGPT, Claude, and Gemini. The transformer model is an encoder-decoder model that encodes an input sequence of text into a sequence of “token”¹ representations then decodes it into predictions of the next position in the sequence by only allowing the decoder to access past positions in the sequence to make its prediction. This is done by “masking” the input sequence in a way that tokens later in the sequence are zeroed out for any given present token’s prediction. At inference time the decoder behaves auto-regressively generating one token in the sequence at a time and using each new prediction as input to generate the next.

The main innovation of the transformer model is its use of self-attention, where it learns a relationship between every pair of tokens in the sequence and then uses that relationship to update its representation of the tokens. Another important part of the transformer architecture is its use of positional encodings. Unlike past language model architectures, like RNN’s or CNN’s, there is no concept of order or position in each token. In order to inject order in transformers order information like some projection of the word’s position is added to each token.

The original transformer model is used to predict the next word and has been adapted to be useful many tasks, the most prevalent being the “helpful assistants” we know today as ChatGPT, Claude, and Gemini.

BERT

Bidirectional Encoder Representations from Transformers (BERT) models are almost identical to the encoder from the original transformer. The difference is instead of doing next word prediction, they are trained with a masked language modeling objective.

How this works is during training 15% of the tokens are “masked” out - changed from their actual word or sub-word representation to a custom “mask” token. As BERT is a transformer, self-attention enables the model to learn a representation for each mask token based on the learned relationships between the mask token and every other token. The model then is trained to predict the true token using all of the surrounding tokens to make its prediction instead of only preceding ones in the sequence. The idea behind BERT is that by being “pre-trained” to predict these “missing” tokens using surrounding tokens as context, the model is “forced” to learn complex patterns between tokens that underlie the language it is being fed. Then, the model can excel at more complex tasks, like sentiment analysis of a movie review or question answering.

ESM-2 Architecture

The ESM-2 architecture and training signal are very similar to BERT. Instead of masking and predicting words and sub-words in sentences, ESM-2 masks and predicts amino acid residues in protein sequences. ESM-2 makes a few modifications to BERT. Namely, it uses Rotary Position Embeddings (RoPE) instead of sinusoidal positional embeddings for its positional encodings². Besides the positional embedding change, ESM-2 also changes hyperparameters like number of layers, number of attention heads, hidden representation dimension size, etc. Multiple sizes of ESM-2 models are trained from the smallest at 8 million parameters to the largest at 15 billion parameters.

Data

Sequence training data for ESM-2 comes from a pair of clustered protein sequence databases, UniRef50 and UniRef90. Many of the catalogued protein sequences come from a variety of sources, so there are a lot of near-duplicates. The goal of the UniRef databases is to reduce this redundancy by clustering similar sequences. The UniRef50 and UniRef90 databases contain the same sequences, but they cluster their proteins by 50% and 90% similarity respectively and define each cluster by a “representative” sequence. At each step of training a minibatch of representative sequences from different UniRef50 clusters are sampled. These sequences are then mapped to the corresponding cluster in UniRef90 and a protein from each of these clusters is randomly sampled. This allows for a diverse training set and results in the model’s processing more than 60 million unique protein sequences over the course of training.

Training

At each step of training, a minibatch of sequences is sampled, a subset of residues are randomly masked, the sequences are processed by the model, and the masked language modeling objective is applied to the predictions for the masked residues. The training proceeds for 500K steps with a batch size of 2 million tokens (270K steps and 3.2 million tokens for the 15B parameter model).

Results

The results of this “fill in the blank” style training are measured using a metric called perplexity. Perplexity intuitively measures on average how many amino acid types the model is choosing between for each prediction, so lower is better: a perplexity of 1 is perfect accuracy and 20 is as bad as random guessing.

When evaluating the trained ESM-2 on a held out set of protein sequences, the ESM-2 authors found that as you increase the size of the model, the “fill in the blank” accuracy increases (perplexity decreases). The 8M parameter model had a perplexity of 10.45, whereas the 15B one had a perplexity of 6.37. The authors also noted that the performance on proteins that had a lot of evolutionarily related proteins in the training set was as good in the smaller scale models as the larger ones. For proteins without that many homologs, however, the larger models were needed to accurately predict the masked residues. This is promising for the future of protein structure prediction, where traditionally the structure prediction accuracy for a protein with a shallower MSA is worse.

The Rest of ESMFold: The “Folding Head” Module

ESMFold is trained by training a “folding head” on top of the ESM-2 model. This folding head is essentially the AlphaFold2 model, but adapted to no longer use MSA’s. The folding head transforms the representation from ESM-2 into 3D coordinates using a two modules: the folding trunk and the structure module.

The Folding Trunk

The folding trunk module is adapted from the Evoformer in the AlphaFold2 model. It consists of 48 folding blocks.

The Evoformer model takes in two representations, the MSA representation and the pairwise representation. Since ESMFold obviates the need for an MSA, it uses the sequence representation from the ESM-2 model instead. The pairwise representation in AlphaFold2 is a tensor that contains a pairwise “interaction” feature for every pair of residues (L^2 total, where L is the length of the protein sequence). The idea is that encoded in this representation is the relative distance and rotation of each pair of residues in 3D space. ESMFold initializes the pairwise representation with relative position information, where the feature for each pair of residues is how many positions apart they are. The pairwise representation then is further refined using information from sequence representation and as it is processed in the folding trunk.

Folding Block

Folding Block: Processing the sequence representation

In the Evoformer, the MSA representation is updated using biaxial attention over the rows and over the columns. Since the sequence representation is just 1D, normal multihead self-attention (as used in the original transformer paper) is used to update it with a bias coming from the pairwise representation.

Processing the pairwise representation

The pairwise representation is updated by taking concatenating an outer product and outer difference between every pair of residue representations in the sequence representation, projecting the concatenated result and adding it to the initial pairwise representation. The pairwise representation is then processed in the same as in the Evoformer in AlphaFold. First, there is the outgoing and incoming triangle multiplicative updates, where a pair of residues’ (i and j) interaction representation, zij is updated by taking each of the two residue’s interaction representations with other residues, elementwise multiplying each interaction with a common residue and summing all the products and then adding the result to the original zij. Second, there is the incoming and outgoing triangle attention, where self-attention between the zij and zik for all k is computed to updated zij as well as zij and zki for all k.

The sequence and pairwise representations are iteratively refined by processing them through 48 folding blocks.

Structure Module

The ESMFold structure module is essentially identical to the one in AlphaFold2.

Frame Initialization

The structure module initializes for each residue a “local frame”, which is essentially a local coordinate system for the backbone of each residue where the alpha carbon is at the origin and the axes are aligned with the carbonyl carbon and nitrogren. The frame is rotation matrix and translation vector that maps coordinates from the local coordinate system to a global coordinate system. At first every residue’s local frame is set so every residue is located at the global origin. As the frames are processed in the structure module, they are gradually nudged outside the origin to their final position prediction.

Invariant Point Attention

The structure module works at high level by using the frames to update the sequence representation with the Invariant Point Attention (IPA) module. The IPA essentially learns local 3D points for each residue maps them back to global space using the frames, uses the distance between each pair of residues in global space to get pairwise attention scores, uses those to update each residue’s local 3D point prediction, and each 3D prediction is used to update the sequence representation. Information from the pairwise representation is also used in the IPA to update sequence representation. The IPA is called invariant because attention scores are invariant to the global transformations of the query, key, and value points, but the entire IPA is equivariant in the sense that as the local frame inputs change, the sequence representation update changes accordingly.

Backbone Update

The output of the IPA, the updated sequence representation, is then passed to the backbone update module, where it is mapped to frame updates using a learned learned linear projection. The frame update is a rotation matrix and translation vector that are multiplied and added to the frame to update it. The IPA and backbone update are then alternated 8 times and the final frame is converted to 3D coordinates to get the final structure prediction.

Loss

The folding head is trained using the Frame Aligned Point Error (FAPE) loss and the distogram loss. The FAPE loss uses each residue’s frame prediction to map the predicted and true 3D structure coordinates to each and every residue’s local frame and then computes L2 difference between the predicted and true atom positions in each frame and then averages the results. The distogram loss predicts inter-residue distances for each pair of residues and then compares them with ground truth distances.

Training

For training data, ESMFold uses 25K clusters of 325K chains of protein structures from PDB along with around 12 million high confidence structures generated with AlphaFold2. At train time they sample predicted AF2 structures 75% of the time and structures from PDB the other 25%. ESMFold was trained for 125K steps then fine-tuned with a structural violation loss for 25K steps.

Results

The results of ESMFold’s structure prediction are evaluated on two test sets: Cameo, which consists of 194 experimentally determined structures not available in PDB, and the CASP14 test set of 51 structures. ESMFold is evaluated on the test set using the TM-score metric, which gives a score between 0 and 1 that measures how similar the predicted and true structure are.

ESMFold was not as accurate as AlphaFold2, getting a score of 0.83 on Cameo, while AF2 got a 0.88. Furthermore ESMFold got a 0.68 on CASP14, whereas AF2 received a 0.85.

The ESMFold authors also noted a few interesting insights from the results. One, lower perplexity (higher accuracy) in the underlying ESM-2 model correlated with stronger structure prediction performance as measured by the TM-score. Secondly, removing or downsizing the folding head does not degrade folding performance as badly as removing the ESM-2 language model.

Conclusion

Training a language model on protein sequences can allow one to train a pretty good structure prediction model without the expensive MSA generation step. ESMFold ends up being up to 60x faster than AF2 at inference time. While the protein language model setup is faster it is not as accurate. Lastly, the bigger language models, the better performance, the masked language modeling performance and the better the structure prediction performance.

1. While us humans have our sentences split up into separate words, we have found that it is easier to train language models when these words are usually further split up (and sometimes combined) into symbols we call tokens! For example, a rare long word like annoyingly might be split up into two tokens: “annoying” and “ly”.

2. RoPE adds positional information by multiplying the query and key projections of each token embedding by rotation matrix that depends on the tokens position in the sequence.