Esm2

ESM2: Evolutionary Scale Modeling

Overview

ESM2 is a family of encoder-only protein language models from Meta FAIR, trained on 250M+ UniRef50 sequences. Unlike ESM3 (which is generative), ESM2 is discriminative — it cannot generate sequences but excels at:

• Protein embeddings — dense representations for downstream ML tasks
• Zero-shot variant effect scoring — predict mutational fitness without labels
• Contact prediction — residue-residue contact maps from attention heads
• Structure prediction — via ESMFold (uses ESM2 as backbone)

ESM2 is available on HuggingFace with no special SDK or API token required. The same HF API also covers ESM-1b and ESM-1v (variant-specialized) — use AutoTokenizer / AutoModel for any of them interchangeably.

Installation

pip install transformers torch

For faster inference with GPU:

pip install transformers torch accelerate

To use the fair-esm package directly (alternative, gives access to ESM-1v and ESM-IF1):

pip install fair-esm

Model Selection

HuggingFace model ID	Params	Layers	Hidden dim	Use case
facebook/esm2_t6_8M_UR50D	8M	6	320	CPU, fast prototyping
facebook/esm2_t12_35M_UR50D	35M	12	480	CPU-friendly, good quality
facebook/esm2_t30_150M_UR50D	150M	30	640	Balanced
facebook/esm2_t33_650M_UR50D	650M	33	1280	Best default — GPU recommended
facebook/esm2_t36_3B_UR50D	3B	36	2560	High accuracy, needs GPU
facebook/esm2_t48_15B_UR50D	15B	48	5120	Max accuracy, multi-GPU

Start with esm2_t33_650M_UR50D unless compute is constrained.

Core Capabilities

1. Protein Embeddings

Extract per-residue or per-sequence embeddings for downstream tasks (classification, clustering, regression).

from transformers import EsmTokenizer, EsmModel
import torch

model_name = "facebook/esm2_t33_650M_UR50D"
tokenizer = EsmTokenizer.from_pretrained(model_name)
model = EsmModel.from_pretrained(model_name).eval()

sequence = "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVVHSLAKWKRQTLGQHDFSAGEGLYTHMKALRPDEDRLSPLHSVYVDQWDWERVMGDGERQFSTLKSTVEAIWAGIKATEAAVSEEFGLAPFLPDQIHFVHSQELLSRYPDLDAKGRERAIAKDLGAVFLVGIGGKLSDGHRHDVRAPDYDDWSTPSELGHAGLNGDILVWNPVLEDAFELSSMGIRVDADTLKHQLALTGDED"

inputs = tokenizer(sequence, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)

# Per-residue embeddings: shape (seq_len, hidden_dim)
# Slice [1:-1] to remove [CLS] and [EOS] special tokens
per_residue = outputs.last_hidden_state[0, 1:-1]

# Per-sequence embedding: mean pool over residues
per_sequence = per_residue.mean(dim=0)  # shape (hidden_dim,)

print(f"Per-residue: {per_residue.shape}")   # (L, 1280)
print(f"Per-sequence: {per_sequence.shape}") # (1280,)

See references/embeddings.md for batch processing, layer selection, and tips for downstream tasks.

2. Zero-Shot Variant Effect Scoring

Score the effect of single amino acid mutations without any labeled data, using masked marginals. A positive score means the mutation is predicted to be beneficial (or neutral); negative means deleterious.

from transformers import EsmTokenizer, EsmForMaskedLM
import torch
import torch.nn.functional as F

model_name = "facebook/esm2_t33_650M_UR50D"
tokenizer = EsmTokenizer.from_pretrained(model_name)
model = EsmForMaskedLM.from_pretrained(model_name).eval()

def score_variant(sequence, position, wt_aa, mut_aa):
    """
    Score a single-site mutation via masked marginals.
    
    Args:
        sequence: wild-type protein sequence (string)
        position: 0-indexed position of the mutation
        wt_aa: wild-type amino acid (single letter)
        mut_aa: mutant amino acid (single letter)
    
    Returns:
        float: log P(mut) - log P(wt) at the masked position.
               Positive = mutation favored by model.
    """
    masked = list(sequence)
    masked[position] = tokenizer.mask_token
    masked_seq = "".join(masked)

    inputs = tokenizer(masked_seq, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits  # (1, L+2, vocab_size)

    # +1 offset for [CLS] token
    log_probs = F.log_softmax(logits[0, position + 1], dim=-1)
    wt_id = tokenizer.convert_tokens_to_ids(wt_aa)
    mut_id = tokenizer.convert_tokens_to_ids(mut_aa)

    return (log_probs[mut_id] - log_probs[wt_id]).item()

# Example: score A2G mutation in a sequence
seq = "MAKVLGKG"
score = score_variant(seq, position=1, wt_aa="A", mut_aa="G")
print(f"A2G score: {score:.3f}")

See references/variant-scoring.md for scanning all positions, scoring multiple mutants, and comparison with ESM-1v.

3. Per-Residue Token Classification

Label each residue for tasks like secondary structure annotation, signal peptide detection, or disorder prediction. Uses EsmForTokenClassification, which adds a per-residue linear head over ESM2 representations.

from transformers import EsmTokenizer, EsmForTokenClassification
import torch

model_name = "facebook/esm2_t33_650M_UR50D"
tokenizer = EsmTokenizer.from_pretrained(model_name)

# Load a fine-tuned checkpoint for your labeling task, or train from scratch:
model = EsmForTokenClassification.from_pretrained(model_name, num_labels=3).eval()
# num_labels: e.g. 3 for helix/sheet/coil, 2 for signal/non-signal, etc.

sequence = "MKTAYIAKQRQISFVKSHFSRQ"
inputs = tokenizer(sequence, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)

# logits: (1, seq_len+2, num_labels) — includes [CLS] and [EOS] special tokens
logits = outputs.logits[0, 1:-1]           # strip specials → (L, num_labels)
predicted_labels = logits.argmax(dim=-1)   # (L,) — one label per residue
print(predicted_labels)

Well-known fine-tuned checkpoints are available on HuggingFace for secondary structure (3-class and 8-class) and signal peptide annotation. The base ESM2 model requires fine-tuning before its token classification outputs are meaningful.

4. Contact Prediction

ESM2 attention heads encode residue-residue contacts. Extract contact maps directly from the model.

from transformers import EsmTokenizer, EsmForProteinFolding
import torch

# Contact prediction uses the folding model
# For lightweight contact prediction without full folding, use fair-esm:
import esm

model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
batch_converter = alphabet.get_batch_converter()
model.eval()

data = [("protein1", "MKTAYIAKQRQISFVK")]
batch_labels, batch_strs, batch_tokens = batch_converter(data)

with torch.no_grad():
    results = model(batch_tokens, repr_layers=[33], return_contacts=True)

contacts = results["contacts"]  # (batch, L, L) — symmetric contact probability matrix
print(contacts.shape)

5. Structure Prediction via ESMFold

ESMFold wraps ESM2 for end-to-end structure prediction. It's fast (seconds per protein) compared to AlphaFold2 and requires only sequence input.

from transformers import EsmForProteinFolding, EsmTokenizer
import torch

tokenizer = EsmTokenizer.from_pretrained("facebook/esmfold_v1")
model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=True)
model.eval()

sequence = "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGD"
inputs = tokenizer([sequence], return_tensors="pt", add_special_tokens=False)

with torch.no_grad():
    output = model(**inputs)

# Convert to PDB
from transformers.models.esm.openfold_utils.protein import to_pdb, Protein as OFProtein
from transformers.models.esm.openfold_utils.feats import atom14_to_atom37

# Helper to save PDB
def output_to_pdb(model_output):
    final_atom_positions = atom14_to_atom37(model_output["positions"][-1], model_output)
    model_output = {k: v.to("cpu").numpy() for k, v in model_output.items()}
    final_atom_positions = final_atom_positions.cpu().numpy()
    pdbs = []
    for i in range(model_output["aatype"].shape[0]):
        aa = model_output["aatype"][i]
        pred_pos = final_atom_positions[i]
        mask = model_output["atom37_atom_exists"][i]
        resid = model_output["residue_index"][i] + 1
        pdbs.append(to_pdb(OFProtein(
            aatype=aa, atom_positions=pred_pos, atom_mask=mask,
            residue_index=resid, b_factors=np.zeros_like(mask),
            chain_index=model_output.get("chain_index", [None] * len(aa))[i],
        )))
    return pdbs

pdb_strings = output_to_pdb(output)
with open("predicted.pdb", "w") as f:
    f.write(pdb_strings[0])

Note: ESMFold requires ~16GB GPU RAM for typical proteins. For CPU-only or large batches, consider the ESMFold web API at https://esmatlas.com/resources?action=fold.

Best Practices

• Special tokens: always strip [CLS] (index 0) and [EOS] (index -1) from last_hidden_state before using per-residue embeddings
• Sequence length: ESM2 supports up to 1022 residues (max_position_embeddings=1026 minus 2 special tokens minus 2 padding); longer sequences must be chunked or truncated
• Normalization: L2-normalize per-sequence embeddings before computing cosine similarity
• Layer choice: last layer is best for structure-related tasks; middle layers (e.g., layer 20 of 33) can be better for evolutionary/functional tasks — see references/embeddings.md
• Variant scoring at scale: for scanning all positions in a sequence, batch the masked inputs rather than running one forward pass per position

Scripts

Bundled CLI scripts for common workflows — no setup beyond pip install transformers torch:

# Batch embed a FASTA file → .npy or HDF5
python scripts/embed.py sequences.fasta --output embeddings.npy

# Scan all single-site mutations for a sequence
python scripts/scan_variants.py MKTAYIAKQRQISFVK --output scores.csv --heatmap landscape.png

# Score specific named variants
python scripts/scan_variants.py MKTAYIAKQRQISFVK --variants A2G K4R T3S

# Use a smaller model for faster prototyping
python scripts/embed.py sequences.fasta --model facebook/esm2_t12_35M_UR50D --output embeddings.npy

Resources

• HuggingFace models: https://huggingface.co/facebook
• GitHub (fair-esm): https://github.com/facebookresearch/esm
• Paper: Lin et al., Science 2023 — https://doi.org/10.1126/science.ade2574
• ESM Metagenomic Atlas: https://esmatlas.com

References

• references/embeddings.md — batch processing, layer selection, normalization, downstream task tips
• references/variant-scoring.md — full position scanning, multi-mutant scoring, ESM-1v comparison, benchmarks
• scripts/embed.py — batch FASTA → embeddings CLI
• scripts/scan_variants.py — mutation landscape scan and named variant scoring CLI