README.md

---

library_name: transformers
license: mit
tags: []
---


# NOTE
The GitHub with the implementation and requirements.txt can be found [here](https://github.com/Synthyra/FastPLMs.git)

# ESM++
[ESM++](https://github.com/Synthyra/FastPLMs) is a faithful implementation of [ESMC](https://biohub.ai/esm/protein) ([license](https://github.com/Biohub/esm/blob/main/LICENSE.md)) that allows for batching and standard Hugging Face compatibility without requiring the ESM Python package.
The large version corresponds to the 600 million parameter version of ESMC.

This repository includes the Biohub ESM MIT license in `LICENSE`.

## Attention backends

`sdpa` (PyTorch Scaled Dot Product Attention) is the default. The backend is set via `config.attn_backend` before loading.

| Backend | Key | Notes |
| :--- | :--- | :--- |
| PyTorch SDPA | `"sdpa"` | Default. Exact numerics, stable on all hardware. |
| Flash Attention | `"kernels_flash"` | Fastest on Ampere/Hopper GPUs. Requires `pip install kernels` (pre-built, no hours-long compilation). Outputs are not bitwise identical to SDPA due to online softmax reordering; differences are often small but not guaranteed to be inconsequential, so use `"sdpa"` if exact numerics matter. |
| Flex Attention | `"flex"` | Skips padding tokens via block mask for faster variable-length batches. Near-exact numerics. First use compiles a Triton kernel (30-120 s). Best combined with `torch.compile`. |
| Auto | `"auto"` | Picks the best available: `kernels_flash`, then `flex`, then `sdpa`. |

```python

from transformers import AutoConfig, AutoModelForMaskedLM



config = AutoConfig.from_pretrained('Synthyra/ESMplusplus_large', trust_remote_code=True)

config.attn_backend = "flex"  # or "kernels_flash", "sdpa", "auto"

model = AutoModelForMaskedLM.from_pretrained('Synthyra/ESMplusplus_large', config=config, trust_remote_code=True)

```

`torch.compile(model)` is heavily recommended for sustained throughput, especially with Flex Attention.

## Binder Design Regularizer

The FastPLMs binder design tutorial uses the ESM++ model family as the
masked-LM pseudoperplexity regularizer while FastPLMs ESMFold2 experimental
models provide differentiable folding losses and final critics. The verified
EGFR example defaults to `Synthyra/ESMplusplus_6B`; this 600M checkpoint exposes
the same `AutoModelForMaskedLM` API and can be used as a lower-memory
regularizer by editing `FastPLMsBinderDesign.lm_name` in
`cookbook/tutorials/binder_design_fastplms.py`.

Default verified run:

```bash

python cookbook/tutorials/binder_design_fastplms.py \

  --backend local \

  --target-name egfr \

  --binder-sequence '################################################################################################################################' \

  --not-antibody \

  --steps 150 \

  --batch-size 1 \

  --seed 103 \

  --output-dir binder_design_egfr_len128_seed103

```

The verified 6B-regularized result had hero mean iPTM `0.913870`, hero min iPTM
`0.904600`, and all four ESMFold2 hero critics above `0.9`.

See [`docs/binder_design.md`](https://github.com/Synthyra/FastPLMs/blob/main/docs/binder_design.md)
for the complete workflow, output files, metrics, and Modal/local compute
options.

## Use with Hugging Face Transformers
```python

from transformers import AutoModelForMaskedLM

model = AutoModelForMaskedLM.from_pretrained('Synthyra/ESMplusplus_large', trust_remote_code=True)

tokenizer = model.tokenizer



sequences = ['MPRTEIN', 'MSEQWENCE']

tokenized = tokenizer(sequences, padding=True, return_tensors='pt')



# tokenized['labels'] = tokenized['input_ids'].clone() # correctly mask input_ids and set unmasked instances of labels to -100 for MLM training



output = model(**tokenized) # get all hidden states with output_hidden_states=True

print(output.logits.shape) # language modeling logits, (batch_size, seq_len, vocab_size), (2, 11, 64)

print(output.last_hidden_state.shape) # last hidden state of the model, (batch_size, seq_len, hidden_size), (2, 11, 1152)

print(output.loss) # language modeling loss if you passed labels

#print(output.hidden_states) # all hidden states if you passed output_hidden_states=True (in tuple)

```

ESM++ also supports sequence and token level classification tasks like ESM2. Simply pass the number of labels during initialization.

```python

from transformers import AutoModelForSequenceClassification, AutoModelForTokenClassification



model = AutoModelForSequenceClassification.from_pretrained('Synthyra/ESMplusplus_large', num_labels=2, trust_remote_code=True)

logits = model(**tokenized).logits

print(logits.shape) # (batch_size, num_labels), (2, 2)

```

ESM++ weights are fp32 by default. You can load them in fp16 or bf16 like this:
```python

import torch

model = AutoModelForMaskedLM.from_pretrained('Synthyra/ESMplusplus_large', trust_remote_code=True, dtype=torch.float16) # or torch.bfloat16

```

## Experimental test-time training

TTT is disabled by default. Normal ESM++ inference, embeddings, logits, and
`state_dict()` keys are unchanged unless you explicitly call `model.ttt(...)`.
The current implementation is experimental and trains only local LoRA adapters
on the ESMC backbone with masked language modeling on the test protein. It can
help some difficult proteins, but it adds test-time compute and can degrade
already confident predictions.

```python

metrics = model.ttt(

    seq="MSTNPKPQRKTKRNT",

    ttt_config={"steps": 3, "ags": 1, "batch_size": 1},

)

model.ttt_reset()

print(metrics["losses"])

```

## Embed entire datasets with no new code
To embed a list of protein sequences **fast**, just call embed_dataset. Sequences are sorted to reduce padding tokens, so the initial progress bar estimation is usually much longer than the actual time it will take.



Example:

```python

embedding_dict = model.embed_dataset(

    sequences=[

        'MALWMRLLPLLALLALWGPDPAAA', ... # list of protein sequences

    ],

    batch_size=2, # adjust for your GPU memory
    max_len=512, # adjust for your needs

    full_embeddings=False, # if True, no pooling is performed

    embed_dtype=torch.float32, # cast to what dtype you want

    pooling_types=['mean', 'cls'], # more than one pooling type will be concatenated together

    num_workers=0, # if you have many cpu cores, we find that num_workers = 4 is fast for large datasets

    sql=False, # if True, embeddings will be stored in SQLite database

    sql_db_path='embeddings.db',

    save=True, # if True, embeddings will be saved as a .pth file

    save_path='embeddings.pth',

)

# embedding_dict is a dictionary mapping sequences to their embeddings as tensors for .pth or numpy arrays for sql

```


```

model.embed_dataset()

Args:

    sequences: List of protein sequences

    batch_size: Batch size for processing

    max_len: Maximum sequence length

    full_embeddings: Whether to return full residue-wise (True) embeddings or pooled (False)

    pooling_type: Type of pooling ('mean' or 'cls')

    num_workers: Number of workers for data loading, 0 for the main process

    sql: Whether to store embeddings in SQLite database - will be stored in float32

    sql_db_path: Path to SQLite database

    

Returns:

    Dictionary mapping sequences to embeddings, or None if sql=True



Note:

    - If sql=True, embeddings can only be stored in float32

    - sql is ideal if you need to stream a very large dataset for training in real-time

    - save=True is ideal if you can store the entire embedding dictionary in RAM

    - sql will be used if it is True and save is True or False

    - If your sql database or .pth file is already present, they will be scanned first for already embedded sequences

    - Sequences will be truncated to max_len and sorted by length in descending order for faster processing

```

## Fine-tuning with Hugging Face PEFT
```python

model = AutoModelForSequenceClassification.from_pretrained('Synthyra/ESMplusplus_large', num_labels=2, trust_remote_code=True)

# these modules handle ESM++ and ESM2 attention layers

target_modules = ["layernorm_qkv.1", "out_proj", "query", "key", "value", "dense"]



lora_config = LoraConfig(

    r=8, # choose lora parameters to your liking

    lora_alpha=16,

    lora_dropout=0.01,

    bias="none",

    target_modules=target_modules,

)



# Apply LoRA to the model

model = get_peft_model(model, lora_config)



# Unfreeze the classifier head

for param in model.classifier.parameters():

    param.requires_grad = True

```

For a more thorough example of fine-tuning, check out our example script [here](https://github.com/Synthyra/FastPLMs/blob/main/fine_tuning_example.py).


## Returning attention maps
When `attn_backend="flex"`, Flex Attention with a pad-token block mask is used for attention calculations. Optimized attention paths do not return attention maps directly.
ESM++ has the option to ```output_attentions```, which will calculate attention manually. This is much slower, so do not use unless you need the attention maps.

```python

output = model(**tokenized, output_attentions=True)

att = output.attentions

len(att) # 33, one for each layer, size (batch_size, num_heads, seq_len, seq_len) each

```

## Comparison across floating-point precision and implementations
We measured the difference of the last hidden states of the fp32 weights vs. fp16 or bf16. We find that the fp16 is closer to the fp32 outputs, so we recommend loading in fp16.
Please note that the ESM package also loads ESMC in fp32 but casts to bf16 by default, which has its share of advantages and disadvantages in inference / training - so load whichever you like for half precision.

Average MSE for FP16: 0.00000003

Average MSE for BF16: 0.00000122

We also measured the difference between the outputs of ESM++ vs. ESMC (both in bfloat16) on 1000 random sequences to ensure compliance with the ESM package.

Average MSE of last hidden state: 2.46e-09

You can load the weights from the ESM package instead of transformers by replacing .from_pretrained(...) to .from_pretrained_esm('esmc_600m')

## Model probes
We employ linear probing techniques on various PLMs and standard datasets, similar our previous [paper](https://www.biorxiv.org/content/10.1101/2024.07.30.605924v1), to assess the intrinsic correlation between pooled hidden states and valuable properties. ESMC (and thus ESM++) perform very well.

The plot below showcases performance normalized between the negative control (random vector embeddings) and the best performer. Classification task scores are averaged between MCC and F1 (or F1max for multilabel) and regression tasks are averaged between Spearman rho and R2.
![image/png](https://huggingface.co/static-proxy/cdn-uploads.huggingface.co/production/uploads/62f2bd3bdb7cbd214b658c48/uRAHYQcwkbgajylTIFbUb.png)

## Inference speeds
We look at various ESM models and their throughput on an H100. Adding efficient batching between ESMC and ESM++ significantly improves the throughput, although ESM++ is also faster than ESMC for batch size one. ESM++ small is even faster than ESM2-35M with long sequences! The most gains will be seen with PyTorch > 2.5 on linux machines.
![image/png](https://huggingface.co/static-proxy/cdn-uploads.huggingface.co/production/uploads/62f2bd3bdb7cbd214b658c48/Lu6nWB9Fc-7YTql3Z1hVB.png)

### Citations

```bibtex

@misc{FastPLMs,

  author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.},

  title={FastPLMs: Fast, efficient, protein language model inference from Hugging Face AutoModel.},

  year={2024},

  url={https://huggingface.co/Synthyra/ESMplusplus_large},

  DOI={10.57967/hf/3726},

  publisher={Hugging Face}

}

```

```bibtex

@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://biohub.ai/papers/esm_protein.pdf},

  note   = {Preprint}

}

```

```bibtex

@article{dong2024flexattention,

  title={Flex Attention: A Programming Model for Generating Optimized Attention Kernels},

  author={Dong, Juechu and Feng, Boyuan and Guessous, Driss and Liang, Yanbo and He, Horace},

  journal={arXiv preprint arXiv:2412.05496},

  year={2024}

}

```

```bibtex

@inproceedings{paszke2019pytorch,

  title={PyTorch: An Imperative Style, High-Performance Deep Learning Library},

  author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith},

  booktitle={Advances in Neural Information Processing Systems 32},

  year={2019}

}

```