Replication of results with the original model

[pengzhangzhi]

Hi,

Report a significant discrepancy between the outputs of your ESM2 implementation (loaded using esme library) and the Hugging Face implementation (facebook/esm2_t6_8M_UR50D). While comparing the logits and representations from both models, I observed sufficiently large differences, which I believe warrant further investigation.

•	Logits:
•	Max Absolute Difference: 1.2461
•	Mean Absolute Difference: 0.1458

•	Representations:
•	Max Absolute Difference: 0.1914
•	Mean Absolute Difference: 0.0226

Is this expected? Or something seems wrong

Below is the code to reproduce the benchmarking

import torch
from esme import ESM2
from esme.alphabet import tokenize
from transformers import EsmTokenizer, EsmForMaskedLM

# ------------------------- Documentation -------------------------
# This script compares the outputs of two implementations of the ESM2 model:
# 1. A custom implementation (loaded using ESM2 and esme.alphabet)
# 2. A pre-trained Hugging Face implementation (using transformers library)
# The comparison focuses on:
# - Logits (model output before applying softmax)
# - Representations (final hidden states of the model)
# It computes and prints the absolute maximum and mean differences between the outputs
# for debugging purposes.

# ---------------------- Load Models and Tokenizer ----------------------
# Load the tokenizer and Hugging Face pre-trained ESM2 model
tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D")
esm = EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D")

# Set the device to GPU (if available) or CPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
esm = esm.to(device)
esm.eval()  # Set the model to evaluation mode

# Load the custom ESM2 model
model = ESM2.from_pretrained("8M.safetensors", device=0)

# ------------------------- Prepare Input Tokens -------------------------
# Define input protein sequences for comparison
sequences = ['MEES', 'M']  # Example sequences
tokens = tokenize(sequences).to(0)  # Tokenize using custom tokenizer

# ---------------------- Custom Model Forward Pass ----------------------
# Generate logits and representations using the custom model
logits_custom = model(tokens).float()  # Model logits
repr_custom = model.forward_representation(tokens).float()  # Model representations

# ---------------------- Hugging Face Model Forward Pass ----------------------
# Generate logits and representations using the Hugging Face model
attention_mask = tokens != tokenizer.pad_token_id  # Attention mask for valid tokens
esm_output = esm(input_ids=tokens, attention_mask=attention_mask, output_hidden_states=True)
logits_esm = esm_output.logits.float()  # Model logits
repr_esm = esm_output.hidden_states[-1].float()  # Last hidden state as representation

# ------------------------- Mask Padding Tokens -------------------------
# Mask out padding tokens for comparison
logits_custom = logits_custom[attention_mask]
logits_esm = logits_esm[attention_mask]
repr_custom = repr_custom[attention_mask]
repr_esm = repr_esm[attention_mask]

# ------------------------- Compute Differences -------------------------
# Compute absolute differences between logits
logit_diff_max = (logits_custom - logits_esm).abs().max()
logit_diff_mean = (logits_custom - logits_esm).abs().mean()

# Compute absolute differences between representations
repr_diff_max = (repr_custom - repr_esm).abs().max()
repr_diff_mean = (repr_custom - repr_esm).abs().mean()

# ------------------------- Print Debug Information -------------------------
# Print logit comparison results
print("Logit Comparison:")
print(f"Max Absolute Difference: {logit_diff_max}")
print(f"Mean Absolute Difference: {logit_diff_mean}")
print(f"Logit Difference Shape: {logits_custom.shape}")

# Print representation comparison results
print("\nRepresentation Comparison:")
print(f"Max Absolute Difference: {repr_diff_max}")
print(f"Mean Absolute Difference: {repr_diff_mean}")
print(f"Representation Difference Shape: {repr_custom.shape}")
eddings.

[MuhammedHasan]

I have comprehensive test cases comparing the original implementation from the GitHub repo of the original ESM model and our implementation. See the following test-cases:

esm-efficient/tests/test_esm2.py

Line 31 in 116076d

def test_esm2_p53(token_p53, flash_esm2, esm2_model):

    p53_human = 'MEEPQSDPSVEPPLSQETFSDLWKLLPENNVLSPLPSQAMDDLMLSPDDIEQWFTEDPGPDEAPRMPEAAPPVAPAPAAPTPAAPAPAPSWPLSSSVPSQKTYQGSYGFRLGFLHSGTAKSVTCTYSPALNKMFCQLAKTCPVQLWVDSTPPPGTRVRAMAIYKQSQHMTEVVRRCPHHERCSDSDGLAPPQHLIRVEGNLRVEYLDDRNTFRHSVVVPYEPPEVGSDCTTIHYNYMCNSSCMGGMNRRPILTIITLEDSSGNLLGRNSFEVRVCACPGRDRRTEEENLRKKGEPHHELPPGSTKRALPNNTSSSPQPKKKPLDGEYFTLQIRGRERFEMFRELNEALELKDAQAGKEPGGSRAHSSHLKSKKGQSTSRHKKLMFKTEGPDSD'

    tokens = tokenize(sequences).to(0)
   
    flash_logit = flash_esm2(token_p53, pad_output=True).float()
    logit = esm2_model(token_p53)['logits']

    perplexity = Perplexity(ignore_index=1).to(device)
    perp_flash = perplexity(flash_logit, token_p53)

    perplexity = Perplexity(ignore_index=1).to(device)
    perp = perplexity(logit, token_p53)
    assert perp > perp_flash

    prob_flash = torch.softmax(flash_logit, dim=-1)
    prob = torch.softmax(logit, dim=-1).float()

    indices_flash = torch.argmax(prob_flash, dim=-1)
    indices = torch.argmax(prob, dim=-1)

    assert torch.all((indices_flash == token_p53) | (indices != token_p53))

    sim = torch.nn.functional.cosine_similarity(
        rearrange(flash_logit, 'b s e -> (b s) e'),
        rearrange(logit, 'b s e -> (b s) e'),
    )
    assert all(sim > .99)

The perplexity of the prediction of esm-efficient is similar to the original ESM model. Moreover, representations have over 99% cosine similarity. Moreover, we have results demonstrating replication of results in the ESM paper.

There might be some absolute difference, but it results from numeral stability during the operations. But does not impact model performance.

Let me know if you have further questions about the results.

[MuhammedHasan]

I also note that mathematically, operations are the same. There will always be some absolute difference due to numeral stability. For example, a @ b is mathematically equallent to (b.T @ a.T).T, but when you run it with PyTorch:

a = torch.rand(2048, 1024).bfloat16() 
b = torch.rand(1024, 4096).bfloat16() 

((a @ b) - (b.T @ a.T).T).abs().max()
... tensor(2., dtype=torch.bfloat16)

This is true for simple matrix multiplication, so consider how many layers and parameters the ESM model has. There will be some numeric differences, but based on the test cases, perplexity calculation, and other results I presented in the paper, these differences are negligible.

[pengzhangzhi]

Hey, thanks! Quick question,

• is ur checkpoint renamed from the official ckpt or u trained urself?
• I have another implementation of ESM2 where only replacing the attention calc with F.scaled_dot_product_attention This implementation has < 1e-5 numerical diff. with the official implementation under the same ckpt.

[pengzhangzhi]

a = torch.rand(2048, 1024).bfloat16()
b = torch.rand(1024, 4096).bfloat16()

I ran the code from my end and got 0 error, not sure what's wrong...

import torch
a = torch.rand(2048, 1024).bfloat16().to('cuda')
b = torch.rand(1024, 4096).bfloat16().to('cuda')
((a @ b) - (b.T @ a.T).T).abs().max()

# tensor(0., device='cuda:0', dtype=torch.bfloat16)

[MuhammedHasan]

a = torch.rand(2048, 1024).bfloat16()

b = torch.rand(1024, 4096).bfloat16()

I ran the code from my end and got 0 error, not sure what's wrong...

import torch

a = torch.rand(2048, 1024).bfloat16().to('cuda')

b = torch.rand(1024, 4096).bfloat16().to('cuda')

((a @ b) - (b.T @ a.T).T).abs().max()



# tensor(0., device='cuda:0', dtype=torch.bfloat16)

Can you run it couple times? I do get non zero results most of the time.

[MuhammedHasan]

Hey, thanks! Quick question,

• is ur checkpoint renamed from the official ckpt or u trained urself?

• I have another implementation of ESM2 where only replacing the attention calc with F.scaled_dot_product_attention This implementation has < 1e-5 numerical diff. with the official implementation under the same ckpt.

Model weights are same as the original with the original esm.

[MuhammedHasan]

See the implementation of Flash-Attention for more details; the implementation is mathematically identical. I only think I changed the transposition of the matrixes and usage of approximated Gelu from Pytorch. Let me know if you have any further specific questions about the code question:

esm-efficient/esme/attention.py

Line 126 in 116076d

class FlashTransformerLayer(nn.Module):

[pengzhangzhi]

Thanks! I'm working on getting them to have close numerical results. I compared the pytorch scaled_dot_attention and the official ESM attention, their error is below 1e-5. So I smell there is def something wrong

[MuhammedHasan]

Let me know if you think there is a difference in the specific part of the code and if so, fell free to reopen the issue.