Skip to content

CSSB-SNU/esm

 
 

Repository files navigation

A world model of protein biology: ESMC, ESMFold2, & ESM Atlas

ESMC & ESMFold2 PreprintAtlasTutorialsSlack

We are releasing a world model for protein biology: a scientific engine for prediction, design, and discovery. Built on the latest generation of Evolutionary Scale Modeling (ESM), this system learns from the protein sequences produced by evolution and uses that knowledge to represent, map, predict, and design proteins across scales — from atomic interactions to evolutionary relationships spanning billions of years. The system includes three artifacts: ESMC, ESMFold2, and ESM Atlas.

ESMC is a state-of-the-art protein language model that has learned the rules of protein biology from training on billions of protein sequences. ESMC defines a new scaling frontier relative to ESM2, achieving stronger performance in emergent long-range structural understanding as model scale increases.

ESMFold2, built on the ESMC 6B model, is a state-of-the-art structure prediction model that has been validated for the design of protein-protein interactions. ESMFold2 surpasses other models in DockQ pass-rate on Foldbench protein-protein and antibody-antigen complexes, and can be used in single-sequence mode for an order of magnitude speedup in folding.

ESMFold2 is validated in the lab across five therapeutic targets. Inversion of ESMFold2 enables generation of de novo minibinders and antibody-derived scFvs with high hit rates, nanomolar affinities, target specificity, and functional activity. We've released the full protocol from target sequence to ranked binder design in this notebook. For additional details, please refer to the preprint.

The ESM Atlas is a map of 6.8 billion proteins covering the full breadth of life’s biodiversity. ESMFold2’s folding throughput enabled the prediction of more than one billion predicted structures. The Atlas is organized according to the internal world model of ESMC. We make this world model interpretable by training sparse autoencoders (SAEs). SAEs are unsupervised neural networks trained to decompose ESMC internal representations into a sparse set of ~16,000 interpretable features that reveal the functional relationships between proteins that ESMC has learned. Each feature is summarized in natural language with an agentic pipeline that maps features onto known biology from protein databases. We release a collection of SAEs trained on different model scales, layers, and at different levels of granularity. Learn more about how to use the ESM Atlas on the Biohub Platform.

For information on using ESM3, see the ESM3 README.

Table of Contents

ESMC

ESMC is a state-of-the-art protein language model that has learned representations of protein biology from training on billions of protein sequences.

Codebase, model weights, and model variants for ESMC are available through Hugging Face.

There are two primary ways of running the ESM models: through the Biohub Platform or locally with Hugging Face. The Biohub Platform enables users to easily run inference with ESM models with minimal setup. Users interested in customizing or fine-tuning ESM models can use the models from Hugging Face.

Running ESMC Locally

Install esm from GitHub (a PyPI release is coming soon):

pip install esm@git+https://github.com/Biohub/esm.git@main

The following code demonstrates how to run ESMC locally

import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer
from huggingface_hub import login

# login with your Hugging Face credentials
login()

# example GFP sequence
sequences = ["MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK"]

model = AutoModelForMaskedLM.from_pretrained(
    "biohub/ESMC-6B",
    device_map="auto",
).eval()
tokenizer = AutoTokenizer.from_pretrained("biohub/ESMC-6B")

inputs = tokenizer(sequences, return_tensors="pt", padding=True)
inputs = {k: v.to(model.device) for k, v in inputs.items()}
with torch.inference_mode():
    output = model(**inputs)

By default, the model returns only the final layer representations. To return hidden states from all transformer layers, set:

output = model(**inputs, output_hidden_states=True)

Running ESMC Through the Biohub Platform

The code below shows how to access ESMC using the Biohub Platform. API tokens can be created in the developer console.

Note that our API migrated from forge.evolutionaryscale.ai to biohub.ai, so some code classes reference “Forge”.

To get started with ESM, install the python library using pip:

pip install esm@git+https://github.com/Biohub/esm.git@main

Then import the necessary libraries and instantiate your desired model.

from esm.sdk import esmc_client
from esm.sdk.api import ESMProtein, LogitsConfig

# Human carbonic anhydrase II (PDB 2CBA)
protein = ESMProtein(
    sequence=(
        "MSHHWGYGKHNGPEHWHKDFPIAKGERQSPVDIDTHTAKYDPSLKPLSVSYDQATSLRILNNGHAFNVEFDD"
        "SQDKAVLKGGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVL"
        "GIFLKVGSAKPGLQKVVDVLDSIKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPPLLECVTWIVLKEP"
        "ISVSSEQVLKFRKLNFNGEGEPEELMVDNWRPAQPLKNRQIKASFK"
    )
)
model = esmc_client(
    model="esmc-600m-2024-12", url="https://biohub.ai", token="<your API token>"
)

protein_tensor = model.encode(protein)
logits_output = model.logits(
    protein_tensor, LogitsConfig(sequence=True, return_embeddings=True)
)

print(logits_output.logits, logits_output.embeddings)

For tutorials on how to use ESMC, see our tutorials.

ESMC Sparse Autoencoders (SAE)

Sparse autoencoders (SAE) are an unsupervised method for decomposing representations of large transformer language models into interpretable units. We released SAEs trained on ESMC to reveal the interpretable units of functional organization that ESMC's world model has learned.

The sparse autoencoder used in the Atlas and analyzed in the paper, ESMC-6B-sae-layer60-k64-codebook16384, is built on the ESMC 6B model. We also provide human-interpretable, agent-generated feature descriptions for this SAE's codebook.

Codebase, model weights, and model variants for ESMC SAEs are available through Hugging Face.

import torch
from transformers import AutoModel, AutoTokenizer

sequence = "MGSNKSKPKDASQRRRSLEPAENVHGAGGGAFPASQTPSKPASADGHRGPSAAFAPAAAEPKLFGGFNSSDTVTSPQRAGPLAGGVTTFVALYDYESRTETDLSFKKGERLQIVNNTEGDWWLAHSLSTGQTGYIPSNYVAPSDSIQAEEWYFGKITRRESERLLLNAENPRGTFLVRESETTKGAYCLSVSDFDNAKGLNVKHYKIRKLDSGGFYITSRTQFNSLQQLVAYYSKHADGLCHRLTTVCPTSKPQTQGLAKDAWEIPRESLRLEVKLGQGCFGEVWMGTWNGTTRVAIKTLKPGTMSPEAFLQEAQVMKKLRHEKLVQLYAVVSEEPIYIVTEYMSKGSLLDFLKGETGKYLRLPQLVDMAAQIASGMAYVERMNYVHRDLRAANILVGENLVCKVADFGLARLIEDNEYTARQGAKFPIKWTAPEAALYGRFTIKSDVWSFGILLTELTTKGRVPYPGMVNREVLDQVERGYRMPCPPECPESLHDLMCQCWRKEPEERPTFEYLQAFLEDYFTSTEPQYQPGENL"

model = AutoModel.from_pretrained("biohub/ESMC-6B", device_map="auto").eval()
tokenizer = AutoTokenizer.from_pretrained("biohub/ESMC-6B")
sae = AutoModel.from_pretrained(
    "biohub/ESMC-6B-sae-k64-codebook16384",
    allow_patterns=["config.json", "layer_30.safetensors", "layer_60.safetensors"],
    device=model.device,
)
sae.initialize_layers([30, 60])
model.add_sae_models([sae.layers["30"], sae.layers["60"]])

inputs = tokenizer(sequence, return_tensors="pt", padding=True)
inputs = {k: v.to(model.device) for k, v in inputs.items()}

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

output["sae_outputs"]["layer60"]  # sparse.coo tensor
print(output["sae_outputs"]["layer60"].shape)

For tutorials on how to use ESMC SAEs, see our tutorials.

ESMFold2

ESMFold2 is a state-of-the-art protein structure prediction model that combines ESMC (6B parameter) language model embeddings with a diffusion-based structure prediction architecture.

The model predicts high-resolution, all-atom 3D protein structures directly from amino acid sequences, with optional multiple sequence alignment (MSA) input for enhanced accuracy on challenging targets. ESMFold2 achieves state-of-the-art performance matching or exceeding AlphaFold3 across diverse evaluation datasets, while offering improved computational efficiency through optimized diffusion sampling and architectural innovations.

Codebase, model weights, and model variants for ESMFold2 are available through Hugging Face

Running ESMFold2 Locally

from esm.models.esmfold2 import (
    DNAInput,
    ESMFold2InputBuilder,
    LigandInput,
    Modification,
    ProteinInput,
    StructurePredictionInput,
)
from transformers.models.esmfold2.modeling_esmfold2 import ESMFold2Model

HHAI_SEQ = (
    "MIEIKDKQLTGLRFIDLFAGLGGFRLALESCGAECVYSNEWDKYAQEVYEMNFGEKPEGDITQVNEKTIPDH"
    "DILCAGFPCQAFSISGKQKGFEDSRGTLFFDIARIVREKKPKVVFMENVKNFASHDNGNTLEVVKNTMNELD"
    "YSFHAKVLNALDYGIPQKRERIYMICFRNDLNIQNFQFPKPFELNTFVKDLLLPDSEVEHLVIDRKDLVMTN"
    "QEIEQTTPKTVRLGIVGKGGQGERIYSTRGIAITLSAYGGGIFAKTGGYLVNGKTRKLHPRECARVMGYPDS"
    "YKVHPSTSQAYKQFGNSVVINVLQYIAYNIGSSLNFKPY"
)

model = ESMFold2Model.from_pretrained("biohub/ESMFold2").cuda().eval()

spi = StructurePredictionInput(
    sequences=[
        ProteinInput(id="A", sequence=HHAI_SEQ),
        DNAInput(
            id="B",
            sequence="GATAGCGCTATC",
            modifications=[Modification(position=5, ccd="C36")],
        ),
        DNAInput(
            id="C",
            sequence="TGATAGCGCTATC",
            modifications=[Modification(position=6, ccd="C36")],
        ),
        LigandInput(id="L", ccd=["SAH"]),
    ]
)

result = ESMFold2InputBuilder().fold(
    model, spi, num_loops=20, num_sampling_steps=100, num_diffusion_samples=1, seed=0
)

print(f"pLDDT mean: {float(result.plddt.mean()):.3f}, pTM: {float(result.ptm):.3f}, ipTM: {float(result.iptm):.3f}")

with open("1mht_pred.cif", "w") as f:
    f.write(result.complex.to_mmcif())

AMD ROCm users: use ROCm 6.4 with PyTorch 2.9 or newer.

Running ESMFold2 Through the Biohub Platform

Install the esm Python package

pip install esm@git+https://github.com/Biohub/esm.git@main

Import the necessary libraries.

from esm.sdk.forge import SequenceStructureForgeInferenceClient
from esm.sdk.api import FoldingConfig
from esm.utils.structure.input_builder import ProteinInput, StructurePredictionInput

Call the inference client with the selected model of choice and replace with your token name.

client = SequenceStructureForgeInferenceClient(model="esmfold2-fast-2026-05", url="https://biohub.ai", token="<your API token>")

# Human carbonic anhydrase II (PDB 2CBA)
ca2_sequence = (
    "MSHHWGYGKHNGPEHWHKDFPIAKGERQSPVDIDTHTAKYDPSLKPLSVSYDQATSLRILNNGHAFNVEFDD"
    "SQDKAVLKGGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVL"
    "GIFLKVGSAKPGLQKVVDVLDSIKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPPLLECVTWIVLKEP"
    "ISVSSEQVLKFRKLNFNGEGEPEELMVDNWRPAQPLKNRQIKASFK"
)
ca2_input = StructurePredictionInput(
    sequences=[ProteinInput(id="A", sequence=ca2_sequence)]
)

config = FoldingConfig(
    num_loops=20,
    num_sampling_steps=100
)
result = client.fold_all_atom(ca2_input, config=config)

with open("result.cif", "w") as f:
    f.write(result.complex.to_mmcif())

For tutorials on how to use ESMFold2, see our tutorials.

Frontier Safety

Biohub has established a safety team to assess the benefits and potential risks of our models and tools prior to release, and develop mitigations where necessary. To do this, we follow a structured approach that includes assessing both biosafety and biosecurity risks as well as existing, comparable open-source models and tools. We actively engage with the scientific community, stakeholders and domain experts to advance innovation as well as best practices for responsible development. Risk assessment was conducted for each of the components of this release, including our ESMC Cambrian models, ESMFold2, ESMC SAEs, ESM Atlas, and binder design system.

Informed by our risk assessments, we are releasing the source code and model weights for ESMC 6B, ESMFold2, and ESMC SAEs. We are also releasing our ESM Atlas dataset and binder design system openly. Biohub values open science, and we share our research with the scientific community so that others can evaluate, reproduce, and build upon our work.

Evaluations: Prior to release, we conducted evaluations to inform our understanding of capability uplift for specific misuse-relevant functional tasks. The full details of these evaluations are available in our corresponding paper appendix.

The Biohub Platform: We implement guardrails that detect and restrict the use of keywords and sequences corresponding to controlled pathogens and toxins on our freely accessible platform. For further details regarding these guardrails, please refer to our Biohub Platform Resources page. We recognize there are many legitimate reasons to use AI models to understand and model these sequences and proteins. If you are a researcher whose work is impacted by these guardrails, you can request elevated access to our platform via biohub.ai.

Please follow our Acceptable Use Policy when using the model.

Licenses

These models are available under the MIT license.

Citations

If you use ESM in your work, please cite one of the following:

ESMC, SAEs, and ESMFold2

@misc{candido2026language,
  title  = {Language Modeling Materializes a World Model of Protein Biology},
  author = {Candido, Salvatore and Hayes, Thomas and Derry, Alexander and Rao, Roshan
            and Lin, Zeming and Verkuil, Robert and Wu, Bryan and Lee, Jin Sub
            and Bruguera, Elise S. and Keval, Jehan A. and Kopylov, Mykhailo
            and Pak, John E. and Wu, Wesley and Thomas, Neil and Mataraso, Samson
            and Hsu, Alvin and Trotman-Grant, Ashton C. and Fatras, Kilian
            and dos Santos Costa, Allan and Badkundri, Rohil and Ak{\i}n, Halil
            and Oktay, Deniz and Deaton, Jonathan and Montabana, Elizabeth
            and Sitwala, Hrishita and Yu, Yue and Wiggert, Marius
            and Carlin, Dylan Alexander and Goering, Anthony W. and Blazejewski, Tomasz
            and Sandora, McCullen and Hla, Michael and Jia, Tina Z.
            and Kloker, Leon H. and Sofroniew, Nicholas J. and Uehara, Masatoshi
            and Pannu, Jassi and Bachas, Sharrol and Liu, Daniel S.
            and Sercu, Tom and Rives, Alexander},
  year   = {2026},
  url    = {https://www.biorxiv.org/content/10.64898/2026.06.03.729735},
  note   = {Preprint}
}

ESM3

@article {hayes2024simulating,
	author = {Hayes, Thomas and Rao, Roshan and Akin, Halil and Sofroniew, Nicholas J. and Oktay, Deniz and Lin, Zeming and Verkuil, Robert and Tran, Vincent Q. and Deaton, Jonathan and Wiggert, Marius and Badkundri, Rohil and Shafkat, Irhum and Gong, Jun and Derry, Alexander and Molina, Raul S. and Thomas, Neil and Khan, Yousuf A. and Mishra, Chetan and Kim, Carolyn and Bartie, Liam J. and Nemeth, Matthew and Hsu, Patrick D. and Sercu, Tom and Candido, Salvatore and Rives, Alexander},
	title = {Simulating 500 million years of evolution with a language model},
	year = {2025},
	doi = {10.1126/science.ads0018},
	URL = {http://dx.doi.org/10.1126/science.ads0018},
	journal = {Science}
}

ESM Github (Code / Weights)

@software{evolutionaryscale_2024,
  author = {{EvolutionaryScale Team}},
  title = {evolutionaryscale/esm},
  year = {2024},
  publisher = {Zenodo},
  doi = {10.5281/zenodo.14219303},
  URL = {https://doi.org/10.5281/zenodo.14219303}
}

ESMFold2 folding & profiling scripts (this fork)

This fork adds two scripts under scripts/:

  • run_esmfold2.py — fold a protein from a JSON job spec (sequence(s) ± MSA), writing .cif + confidence (pTM/PAE/pLDDT).
  • profile_esmfold2.py — profile inference: per-module time & peak GPU memory, per-block deep dives, Perfetto GPU traces, and TensorBoard profiles.

Build / install

Environment: a uv-managed venv on Python 3.12. On the cluster, load the CUDA toolkit used for any source builds:

module load cuda/12.8          # provides nvcc 12.8 at /opt/ohpc/pub/apps/cuda/12.8

Always target the project venv explicitly (the login shell auto-activates a shared anaconda; plain uv pip would hit that): VIRTUAL_ENV=.venv uv pip install --python .venv/bin/python ...

Required

  1. PyTorch 2.8.0 + cu126 (pinned). torch.compile / apply_torch_compile for ESMFold2 fails on torch 2.12 (Inductor unbacked-symbol error in the atom unpadding); 2.8.0 is the version Biohub's CI validates.
    VIRTUAL_ENV=.venv uv pip install --python .venv/bin/python \
      --index-url https://download.pytorch.org/whl/cu126 \
      --extra-index-url https://pypi.org/simple \
      --index-strategy unsafe-best-match  torch==2.8.0+cu126
  2. Biohub transformers fork (ships transformers/models/esmfold2) — pinned in pyproject.toml (transformers @ git+https://github.com/Biohub/transformers).
  3. flash-attnrequired: transformers/models/esmc imports it unguarded, so ESMC won't load without it. No prebuilt wheel exists for this torch; build from source against torch 2.8 with a clean (non-conda) gcc toolchain:
    srun -p cpu --mem=320G --cpus-per-task=16 bash -c '
      unset CONDA_PREFIX CPATH CPLUS_INCLUDE_PATH CPPFLAGS DEBUG_CPPFLAGS   # de-conda the build
      export PATH=$(echo "$PATH" | tr ":" "\n" | grep -v anaconda3 | paste -sd:)
      export CUDA_HOME=/opt/ohpc/pub/apps/cuda/12.8 PATH=$CUDA_HOME/bin:$PATH
      export CC=/opt/ohpc/pub/compiler/gcc/9.4.0/bin/gcc CXX=/opt/ohpc/pub/compiler/gcc/9.4.0/bin/g++
      export MAX_JOBS=8 TORCH_CUDA_ARCH_LIST=8.6 FLASH_ATTENTION_FORCE_BUILD=TRUE
      VIRTUAL_ENV=.venv uv pip install --python .venv/bin/python flash-attn --no-build-isolation --no-deps'
    (~45 min; high RAM at MAX_JOBS=8. Rebuild whenever torch changes — it's ABI-pinned.)

Strongly recommended — cuEquivariance kernels (the large speedup: ~7–8× trimul):

VIRTUAL_ENV=.venv uv pip install --python .venv/bin/python cuequivariance-torch
VIRTUAL_ENV=.venv uv pip install --python .venv/bin/python \
  --extra-index-url https://pypi.nvidia.com cuequivariance-ops-torch-cu12

load_model auto-selects cueq when installed; _enable_cueq_ops() preloads the ops' libcue_ops.so so no manual LD_LIBRARY_PATH is needed.

Optional — profiling viz (for profile_esmfold2.py --tensorboard):

VIRTUAL_ENV=.venv uv pip install --python .venv/bin/python tensorboard torch-tb-profiler

Run

srun -p gpu --gres=gpu:A6000:1 --mem=96G --cpus-per-task=8 \
  .venv/bin/python scripts/run_esmfold2.py --input job.json --dtype float32

Defaults to the most efficient path (cueq on trunk and MSA trimul + OPM chunking

  • TF32). Add --offload-lm for long sequences / 24 GB GPUs. float32 is the supported dtype (bf16 has a diffusion LayerNorm dtype issue).

Our contributions

Optimizations added on top of the stock model (built-in pieces we enabled are noted as such):

  • cueq-MSA hybrid (new) — the upstream set_kernel_backend never reaches the MSA encoder, so its TriangleMultiplicativeUpdates ran pure-PyTorch. We wire cueq into them → MSA encoder 5–8× faster, lower peak memory. Now the default.
  • MSA OuterProductMean chunking (enabled an unwired built-in) — OPM's [L,L,d_hidden²] outer product is the largest transient; the chunked path existed but was never wired to the MSA encoder. Enabling it lets large complexes fit (e.g. H2340 @ 2046 res: OOM → runs on a 48 GB A6000).
  • ESM-C CPU offload (--offload-lm, LMOffloader) (new) — parks the ~24 GB ESMC backbone on CPU during the trunk/diffusion phases; required for long sequences / smaller GPUs.
  • TF32 on the fp32 diffusion matmuls (new) — matches the paper's inference config (the code didn't set it).
  • _enable_cueq_ops() preload (new) — makes the cueq ops loadable without LD_LIBRARY_PATH.
  • Standard ESMFold2Model default — switched off the deprecated experimental arch and fixed a redundant double-ESMC-load.
  • profile_esmfold2.py (new) — non-invasive profiler: per-module time/peak memory, --deep per-block, --mode {cueq,compile,hybrid,cueq-msa}, --torch-trace (Perfetto GPU+memory), --tensorboard (+ allocator memory snapshot), --offload-lm, --kernel-backend, --trunk-layers/--msa-layers, --opm-chunk.

Key benchmark results

All on a single NVIDIA A6000 (48 GB), float32, outputs unchanged (pTM/pLDDT within run-to-run noise).

End-to-end forward (standard model, 16 recycles, --offload-lm):

target residues stock cueq + cueq-MSA + OPM-chunk (ours) speedup
H2343 773 123.8 s 105.5 s 1.17×
H1382 1378 469.8 s 284.8 s 1.65×

Per-component (the levers behind the above):

component stock optimized speedup
trunk trimul — cuEquivariance kernel (enabled) 557 ms/block 72.5 ms/block 7.7×
MSA encoder — cueq-MSA (ours), H2343 2.45 s/iter 0.50 s/iter 4.9×
MSA encoder — cueq-MSA (ours), H1382 12.45 s/iter 1.48 s/iter 8.4×

Memory / feasibility:

lever effect
ESM-C CPU offload (ours) trunk-phase peak GPU 33.5 → 9.6 GB (frees ~24 GB)
cueq-MSA (ours) H1382 MSA-encoder peak 32.1 → 23.4 GB
MSA OPM chunking (ours) H2340 @ 2046 res: OOM → runs (removes an 8.08 GiB transient)

The trunk speedup comes from the built-in cuEquivariance kernels we enabled; the MSA-encoder, memory, and large-complex feasibility wins are this fork's additions (cueq-MSA wiring, OPM chunking, ESM-C offload, TF32).

About

Biohub ESM github + simple optimization

Resources

License

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages

  • Jupyter Notebook 76.3%
  • Python 23.5%
  • Other 0.2%