The 3 type of AI models in protein biology

There are 3 types of AI models that are predominantly used in Biology. Like most other progress in the field of AI, they follow what was already prevalent in the field of natural language understanding.

1. Encoder only models

Best example - ESM-2 from Meta.

Encoder only models use encoder architecture from transformers to learn the vector representation of a protein. A vector representation is super important because vector is a mathematical object and mathematical modeling can be done with it. Without a vector representation it is impossible to do any machine learning. 

There are few specific properties of these vector embeddings which include:

• They are of a fixed length - regardless of the length of the amino acid sequence. This fixed length representation is also called dense vector representation.

• The individual scaler elements of the vector are real numbers.

• This vector encodes a meaningful interpretation of proteins. 2 proteins with close amino acid profile will have similar vector representation. 

  • But a word of caution - you can't check 2 vectors manually to see if they are close or not. You can do that by calculating cosine similarity or Euclidean distance between 2 vectors. 

  • There is no specific threshold of when 2 vectors are considered close. Vector similarity is relative. 

In ESM-2 model, each amino acid gets its own vector representation. Then the final vector representation of the entire protein is done by mean pooling - taking the average of all the vectors. While this is not the only way to generate protein level embedding, it is the most common. 

The total token library for ESM-2 is 33. It consists of 20 standard amino acids, 8 ambiguous or rare amino acids and 5 special tokens such as <cls>, <eos>, <pad>, <unk>, <mask>.

The ESM-2 follows the encoder block architecture first introduced in "Attention is all you need" paper from Google and closely resembles encoder only model Bert. 

In the encoder type architecture, after the model input, it passes through series of transformer blocks. All of these transformer blocks within the same model are always replicates of each other. This means, the transformer block inside ESM-2 are all replicates but are different from encoder blocks inside Bert.

Another thing to note is that every single token of the sequence passes through the same transformer block in parallel. The transformation weight matrix remains the same for each token in same forward pass - it does get updated during backpropagation and in next forward pass each token uses the updated matrix.

The attention is computed by combination of three matrices called as query, key and value. There is a mask for attention that is also maintained. This mask is binary vector and stores the boolean logic to determine if specific position has to be included in attention calculation or not. This is useful during padding when making the vector lengths same for passing in the same length. Attention is not calculated once in each layer - but more than once and this is called as multi headed attention. The final attention score is an average of all the parallel attention blocks.

ESM-2 also has something unique after transformer block which is called ESM contact map head. Basically in this contact map head, the model tries to predict probabilities that which two amino acids are maybe in potential contact with each other. 

Finally there is the standard Language model head which predicts the logits for each of the position. During the training phase the model is masked and the task is to predict the masked amino acid. The better the model gets at this, the better the overall quality of the model.

One of the most important element inside any neural network model is the choice of loss function. Inside ESM-2 also the loss function used is cross-entropy - a standard for masked language modelling. In MLM, the loss function is calculated only by masked token. Essentially the non masked token are used to predict the logit probability for each masked position. Based on the difference between the actual token and the predicted token for the masked position, one is able to calculate the loss. The goal is to minimize this loss as much as possible.

2. Decoder only models - Progen2

Decoder only models are second part of the transformer architecture. They are considered auto-regressive models because they take predict one token at a time. The next token. Another difference is that the attention calculation is done by looking at only the previous tokens and not the next tokens. This is called as masked attention.

The Weight matrix for key, query and value remain same for all the tokens and are trained during backpropagation.

The loss function is still cross entropy. It is calculated as the difference in the log probability of the predicted token vs the actual token. The model keeps producing something until it hits the <eos> token.

With proteins, there is a N terminal and there is a C terminal. Progen is trained such that it can produce protein from each side. Biologically though the production of protein is done from N terminal to C terminal.

Progen2 is also trained on conditional tags to be able to generate proteins based on specific natural language prompts. IMO, the quality of that is bit questionable. 

3. Diffusion type models (Evo-diff or rfdiffusion)

4. Inverse folding models - ESM-IF/ProteinMPNN: This model takes a structure and then keeps some residue conserved and then predict the sequence