Implementing Protein Sequence Mutation Effect Prediction with ESM Models

[monicamtavares]

Hello everyone,

I am currently working on a project to predict the effects of single mutations in protein sequences using transformer-based models. Specifically, I am interested in utilizing the following models: ESM-1b (esm1b_t33_650M_UR50S), ESM-1v (esm1v_t33_650M_UR90S_[1-5]), and ESM-2 (esm2_t33_650M_UR50D).

My aim is to extract what has been termed the "mutation probability" prediction from these models, which represents the model's estimation of the likelihood of various mutations at each position in the sequence. I have written some code to do this, and I would be very grateful for any input on whether I am correctly implementing this process. Furthermore, I would like to know if the same implementation would work across all three models.

Here is my current approach:

import torch
import esm

# Load ESM-1b model
model, alphabet = esm.pretrained.esm1b_t33_650M_UR50S()
batch_converter = alphabet.get_batch_converter()
model.eval()  # disable dropout for deterministic results

# Use CUDA if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

from scipy.special import softmax

def extract_features(model, batch_converter, device, df, seq_column):
    
    #Extracts mutation probability features from a DataFrame of sequences.
    
    # Prepare to store the results
    all_probs = []
    
    # Iterate over the DataFrame
    for i, sequence in enumerate(df[seq_column]):
        
        # Prepare data
        data = [("protein" + str(i), sequence)]
    
        # Convert a batch of sequences into Torch tensors using batch_converter
        batch_labels, batch_strs, batch_tokens = batch_converter(data)
        batch_tokens = batch_tokens.to(device)

        # Extract logits
        # Disable contacts and repr_layers to only get mutation predictions
        with torch.no_grad():
            results = model(batch_tokens, repr_layers=[], return_contacts=False)
        # Get the logits
        logits_np = (results['logits'].detach().cpu().numpy()[0,].T)[4:24,1:-1]
        # Apply scipy's softmax function
        all_probs.append(softmax(logits_np, axis=-1))
        
        # Clean up
        del batch_tokens, results
        torch.cuda.empty_cache()
        
    return all_probs

Any feedback on my code or approach would be greatly appreciated. Thank you in advance!

[gieses]

I just have a couple of comments from my own experiments & recent papers

• you can get pre-computed scores for ESM1b from this paper: https://www.biorxiv.org/content/10.1101/2022.08.25.505311v1.full
• naively, you can only compute scores for proteins <= 1024 length with ESM1*. The above paper shows that you can do a relatively simple sliding window to get the scores for longer proteins. ESM2 doesnt have that problem
• from my experiments ESM2 (650M, 3B) wasnt better than ESM1b but similarly good.
• I think in the original implementations for variant scoring, you need to mask each residue individually to retrieve the scores. In a recent paper from Kevin Yang and co I saw a plot that shows that correlation of logits of masked vs. non-masked residues is very high. So maybe you do not need to mask
• if you plan to compute many scores, it's worthwhile parallelizing more, e.g. doing a larger batch of proteins and use torch softmax.

[monicamtavares]

Hi @gieses, I appreciate your suggestions and the reference to the paper; it's indeed a valuable resource. I am, in fact, trying to write code that produces the prediction scores for ESM1b similar to the methodology they used.

The protein that I am currently studying is 403 amino acids long, so it falls within the limit for ESM1b.

As for the masking issue, I understand that in the original implementations for variant scoring, each residue is masked individually to retrieve the scores. I created a function to do this, which you can see below:

def find_and_mask_difference(wild_type_sequence, mutated_sequence):
    assert len(wild_type_sequence) == len(mutated_sequence), "Sequences should be of the same length"
    for i in range(len(wild_type_sequence)):
        if wild_type_sequence[i] != mutated_sequence[i]:
            masked_sequence = wild_type_sequence[:i] + '<mask>' + wild_type_sequence[i+1:]
            return masked_sequence, i
    return wild_type_sequence, -1  # return the original sequence and -1 if no difference found

The challenge I'm facing right now is that my function does not specify the amino acid that was substituted. In the original implementations for variant scoring, there is a score for each possible amino acid substitution at each position. My objective, however, is to predict scores specifically for certain amino acid substitutions at given residues. I would appreciate any suggestions on how to adapt my function to achieve this. Alternatively, should I attempt to address this issue when interpreting the logits results?

[OliverT1]

Hi @gieses - do you have a reference for the Kevin Yang paper comparing masked vs not masking for likelihood estimation?

[gieses]

Hi @OliverT1, I actually went looking for the reference but could not find it, sorry. Hope I didnt mess anything up.
It was only a small plot (even supplement afaik).

[eduardoftoliveira]

I think I know Kevin Yang's paper you are referring to, https://www.biorxiv.org/content/10.1101/2023.03.04.531015v1