Protein Language Models

protein-language-models

When to Use

• Need per-residue or per-sequence embeddings from protein sequence for downstream ML (classification, clustering, similarity search)
• Need fast single-sequence structure prediction without building an MSA (ESMFold)
• Scoring the fitness effect of point mutations zero-shot — deep mutational scanning triage, variant-of-uncertain-significance (VUS) prioritization — with ESM-1v
• Designing or redesigning a protein sequence for a fixed backbone (inverse folding) with ESM-IF1
• Extracting attention-derived contact maps as a fast alternative to full structure prediction

Version Compatibility

• fair-esm ≥ 2.0.0 (pip install fair-esm), PyTorch ≥ 2.0, Python ≥ 3.9
• CUDA GPU recommended for ESM2-650M+ and ESMFold; ESM2-8M/35M run fine on CPU
• biotite ≥ 0.39 required for ESM-IF1 structure loading

Prerequisites

• pip install fair-esm torch biotite matplotlib
• Familiarity with PyTorch tensors and basic protein sequence/FASTA concepts (see bio-sequence-manipulation-seq-objects)
• ESMFold and ESM-IF1 strongly prefer a CUDA GPU with ≥16GB VRAM

Quick Reference

Goal: turn protein sequences into fixed-length vectors for downstream ML (clustering, similarity search, classifiers). Approach: run ESM2, extract per-residue representations from the final layer, mean-pool over the true (unpadded) residue span per sequence.

import esm
import torch
import numpy as np

def embed_sequences(seqs, model_name='esm2_t33_650M_UR50D', layer=33):
    """Compute per-sequence ESM2 embeddings via mean-pooling residue representations.

    Args:
        seqs: list of (name, sequence) tuples
        model_name: attribute name on esm.pretrained (e.g. 'esm2_t33_650M_UR50D')
        layer: transformer layer to pool from (33 = last layer of the 650M model)
    Returns:
        np.ndarray of shape (n_seqs, embed_dim)
    """
    model, alphabet = getattr(esm.pretrained, model_name)()
    model.eval()
    batch_converter = alphabet.get_batch_converter()
    _, _, batch_tokens = batch_converter(seqs)

    with torch.no_grad():
        results = model(batch_tokens, repr_layers=[layer], return_contacts=False)
    token_rep = results['representations'][layer]

    # Slice per true sequence length (not the batch max) to avoid padding tokens
    # in the mean, and skip index 0 (<cls>/BOS token).
    embeddings = [
        token_rep[i, 1:len(seq) + 1].mean(0).numpy()
        for i, (_, seq) in enumerate(seqs)
    ]
    return np.vstack(embeddings)

data = [
    ('protein_A', 'MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGD'),
    ('protein_B', 'MKTLLLTLVVVTIVCLDLGYTPETRFLLKFNKAVIVAGTNTK'),
]
seq_embeddings = embed_sequences(data)
print(seq_embeddings.shape)  # (2, 1280)

Goal: predict a 3D structure and per-residue confidence directly from sequence, no MSA needed. Approach: load ESMFold once, call .infer() for the confidence dict and .infer_pdb() for a writable PDB string.

import esm
import torch

def predict_structure(sequence, out_pdb='esmfold_structure.pdb'):
    """Predict structure with ESMFold, write a PDB, and return mean pLDDT.

    pLDDT interpretation: >90 very confident, 70-90 confident, 50-70 low,
    <50 likely disordered/unreliable.
    """
    model = esm.pretrained.esmfold_v1().eval().to('cuda')
    with torch.no_grad():
        pdb_str = model.infer_pdb(sequence)
        output = model.infer(sequence)
    with open(out_pdb, 'w') as f:
        f.write(pdb_str)
    plddt = output['plddt'].squeeze().cpu().numpy()
    return float(plddt.mean())

mean_plddt = predict_structure('MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEK')
print(f'Mean pLDDT: {mean_plddt:.2f}')

Goal: rank candidate point mutations without any labeled fitness data. Approach: mask the mutated position, compare the model's log-probability of the wild-type vs. mutant residue at that position (masked marginal method).

import esm
import torch

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

def score_mutation(wt_seq, position, wt_aa, mut_aa):
    """Score a mutation's effect via masked marginal log-likelihood.

    Args:
        position: 0-indexed position in wt_seq
    Returns:
        log P(mut | context) - log P(wt | context); negative = destabilizing/deleterious
    """
    _, _, batch_tokens = batch_converter([('protein', wt_seq)])
    masked_tokens = batch_tokens.clone()
    masked_tokens[0, position + 1] = alphabet.mask_idx  # +1 for the <cls>/BOS token

    with torch.no_grad():
        logits = model(masked_tokens)['logits']
    log_probs = torch.nn.functional.log_softmax(logits[0, position + 1], dim=-1)

    wt_idx = alphabet.get_idx(wt_aa)
    mut_idx = alphabet.get_idx(mut_aa)
    return (log_probs[mut_idx] - log_probs[wt_idx]).item()

wt_sequence = 'MDLSALRVEEVQNVINAMQKILECPICLELIKEPVSTKCDHIFCKFCMLKLLNQKKGPSQCPLCKNDITKRSLQESTRFSQLVEELLKIICAFQLDTGLEYANSYNFAKKENNSPEHLKDEVSIIQSMGYRNACKESMLCTHSSLNFFPVSLNLNPFQNNRNQLQNELREQLKLRQLEMDLNRFLSEYRSSMSLNHLENSSAAQLKLMQQKELNQIFQELNFQNQNQNQNQNQ'
delta_llr = score_mutation(wt_sequence, 1, 'D', 'E')  # D2E mutation, position 2 (1-indexed)
print(f'Delta log-likelihood: {delta_llr:.3f}')

Goal: get a fast contact map without running a full structure predictor. Approach: enable return_contacts=True in the same ESM2 forward pass used for embeddings.

import esm
import torch
import matplotlib.pyplot as plt

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

_, _, tokens = batch_converter([('protein', 'MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGD')])
with torch.no_grad():
    results = model(tokens, repr_layers=[33], return_contacts=True)

contacts = results['contacts'][0]  # (L, L) contact probability map
plt.matshow(contacts.numpy(), cmap='RdYlBu_r')
plt.title('ESM2 Contact Prediction')
plt.colorbar()

Goal: generate new sequences that are predicted to fold onto a given backbone (fixed-backbone design). Approach: load coordinates from a PDB with ESM-IF1's utilities, then sample sequences at a chosen temperature.

import esm
import esm.inverse_folding

model, alphabet = esm.pretrained.esm_if1_gvp4_t16_142M_UR50()
model.eval()

structure = esm.inverse_folding.util.load_structure('protein.pdb', 'A')
coords, native_seq = esm.inverse_folding.util.extract_coords_from_structure(structure)

# temperature=1.0 -> diverse samples; temperature=0.1 -> conservative, close to native
sampled_seq = model.sample(coords, temperature=1.0, partial_seq=None)

Pitfalls

• Padding: ESM pads to the batch's max length; mean-pool using each sequence's true length (as above), not the padded batch width — otherwise short sequences get diluted by padding-token representations
• Token offset: ESM prepends <cls> at position 0, so residue i (0-indexed) lives at token index i+1
• ESMFold speed/accuracy: ~0.5s/sequence on GPU vs. minutes-hours for AlphaFold2 with an MSA; ESMFold accuracy is somewhat lower, especially for shallow-MSA/orphan proteins
• Zero-shot scoring: masked-marginal scoring captures single-mutation effects well but does not model epistasis between simultaneous mutations
• Sequence length: ESM2/ESM-1v max context is 1024 tokens; window or truncate longer proteins and average per-window scores
• Model download size: 650M+ checkpoints are multi-GB downloads on first esm.pretrained.*() call; cache the ~/.cache/torch/hub/checkpoints directory across runs

See Also

• ai-science-esm2-embeddings — narrower embedding/ESMFold quickstart using the same library
• alphafold-structure-prediction — MSA-based structure prediction and pLDDT/PAE interpretation
• bio-core-protein-structure — parse/analyze the PDB output of ESMFold or ESM-IF1 (RMSD, DSSP, contacts)
• ai-science-alphafold-protein-design — triage RFdiffusion/ProteinMPNN design candidates alongside ESM-IF1 outputs