Synthyra/FastESMFold

NOTE

The GitHub with the implementation and requirements.txt can be found here

FastESMFold

FastESMFold is a self-contained, HuggingFace-compatible reimplementation of ESMFold with optional Test-Time Training (TTT) and multi-backend attention (SDPA, Flash, Flex).

No dependency on fair-esm, proteinttt, or openfold. Just transformers, torch, and einops.

Why Test-Time Training?

Protein language models like ESM2 are trained on millions of sequences, but at inference time they process each new protein in a single forward pass with no adaptation. This is a missed opportunity: the input sequence itself contains structural signal that the model could learn from.

Test-Time Training (TTT) adapts the model to each individual protein before predicting its structure. The idea is simple: before folding, we briefly train the ESM2 backbone on the input sequence using masked language modeling (the same objective it was pretrained with). This forces the model to "study" the specific sequence, strengthening its internal representation of that protein's structural features.

The adaptation uses LoRA (Low-Rank Adaptation) for efficiency: only small adapter weights are trained (~4.4M parameters out of 3.5B), and the base model is restored after each prediction. This takes 20-45 seconds per sequence on an A10G GPU but can dramatically improve structure prediction quality, especially on difficult targets where standard ESMFold produces low-confidence predictions.

When is TTT most useful?

• Sequences with low baseline pLDDT (< 0.5): TTT can improve pLDDT by 10-30+ points
• Novel proteins with limited homology in training data
• Disordered or multi-domain proteins where ESMFold struggles

When is TTT unnecessary?

• Sequences that already fold well (baseline pLDDT > 0.7): TTT rarely helps and may slightly degrade predictions
• High-throughput screening where speed matters more than accuracy

Key Features

• Standard ESMFold: Full ESMFold v1 structure prediction, loadable via AutoModel
• Optional TTT: Enable test-time training for improved structure prediction on difficult sequences
• Best structure selection: When TTT is enabled, folds after each step and returns the structure with the highest pLDDT
• FastESM2 attention: SDPA/Flash/Flex backends for the 3B ESM2 backbone
• Self-contained LoRA: lora_diffusion-compatible implementation (no peft dependency)
• 3.5B parameters: Full ESMFold v1 architecture (ESM2-3B backbone + folding trunk)

Use with transformers

Standard structure prediction (no TTT)

import torch
from transformers import AutoModel

model = AutoModel.from_pretrained(
    "Synthyra/FastESMFold",
    trust_remote_code=True,
    torch_dtype=torch.float32,
).cuda().eval()

# Standard fold (no TTT)
with torch.no_grad():
    output = model.infer("MKTLLILAVVAAALA...")
pdb_strings = model.output_to_pdb(output)
plddt = output["plddt"].mean().item()
print(f"pLDDT: {plddt:.3f}")

Structure prediction with TTT

TTT adapts the ESM2 backbone to a specific input sequence via masked language modeling before folding. This can dramatically improve pLDDT on difficult sequences (e.g., 0.38 to 0.72).

# Configure TTT
model._ttt_cfg.steps = 10      # 10 optimizer steps (default)
model._ttt_cfg.lora_rank = 8   # LoRA rank (default)
model._ttt_cfg.lora_alpha = 32 # LoRA scale (default)

# fold_protein() runs TTT, folds after each step, returns best structure
result = model.fold_protein("MKTLLILAVVAAALA...")
print(f"pLDDT: {result['plddt']:.3f}")
print(f"Best step: {result['best_step']} (0=baseline, 1-10=TTT steps)")
print(f"Step pLDDTs: {[f'{p:.2f}' for p in result['step_plddts']]}")

# Save PDB
with open("structure.pdb", "w") as f:
    f.write(result["pdb_string"])

Return values

fold_protein(sequence) returns a dict:

Key	Type	Description
plddt	float	Best mean pLDDT across all TTT steps
ptm	float	Predicted TM-score from best step
pdb_string	str	PDB format structure from best step
step_plddts	list[float]	pLDDT at each step [baseline, s1, ..., s10]
best_step	int	Which step produced the best structure (0=baseline)

Disabling TTT

To use FastESMFold as a standard ESMFold (no TTT), set steps=0 or call infer() directly:

# Option 1: Set TTT steps to 0
config = AutoConfig.from_pretrained("Synthyra/FastESMFold", trust_remote_code=True)
config.ttt_config = {"steps": 0}
model = AutoModel.from_pretrained("Synthyra/FastESMFold", config=config, trust_remote_code=True)
result = model.fold_protein("MKTLLILAVVAAALA...")  # No TTT, just baseline fold

# Option 2: Call infer() directly (inherited from EsmForProteinFolding)
with torch.no_grad():
    output = model.infer("MKTLLILAVVAAALA...")
pdb_strings = model.output_to_pdb(output)

TTT Benchmark

Tested on 10 difficult sequences on A10G GPU:

Metric	Value
Mean baseline pLDDT	0.549
Mean best TTT pLDDT	0.637
Mean improvement	+0.088
Sequences improved >5pt	5/10
Time per sequence	~20-45s
GPU memory peak	18.3 GB

On the hardest sequence (baseline pLDDT 0.38), TTT improves to 0.72 (+34 points).

Attention backends

The ESM2 backbone supports multiple attention backends via config.attn_backend:

Backend	Key	Notes
PyTorch SDPA	"sdpa"	Default. Exact numerics, stable on all hardware.
Flash Attention	"kernels_flash"	Fastest. Requires pip install kernels.
Flex Attention	"flex"	Skips padding tokens via block mask. First use compiles a Triton kernel.
Auto	"auto"	Picks best available: kernels_flash > flex > sdpa.

from transformers import AutoConfig, AutoModel

config = AutoConfig.from_pretrained("Synthyra/FastESMFold", trust_remote_code=True)
config.attn_backend = "kernels_flash"
model = AutoModel.from_pretrained("Synthyra/FastESMFold", config=config, trust_remote_code=True)

TTT Configuration

TTT parameters are set via config.ttt_config (a dict) or by modifying model._ttt_cfg after loading:

Parameter	Default	Description
lr	4e-4	Learning rate for SGD optimizer
steps	10	Number of optimizer steps (0 to disable TTT)
ags	4	Gradient accumulation steps per optimizer step
batch_size	4	Batch size for masked language model training
mask_ratio	0.15	Fraction of tokens to mask
lora_rank	8	LoRA rank (0 for full backbone fine-tuning)
lora_alpha	32.0	LoRA scaling factor (applied as scale=alpha, matching lora_diffusion)
seed	0	Random seed for reproducible LoRA initialization and masking
lora_target_class	"EsmSelfAttention"	Which module class to inject LoRA into

How TTT Works

1. Baseline fold (step 0): Standard ESMFold prediction
2. LoRA injection: Rank-8 LoRA adapters on all nn.Linear layers inside ESM2 attention modules
3. Masked LM training: 10 optimizer steps (each with 4 gradient accumulation sub-steps) of BERT-style masked language modeling on the input sequence
4. Per-step folding: After each optimizer step, fold the sequence and record pLDDT
5. Best selection: Return the structure with the highest pLDDT
6. Reset: Restore LoRA weights to initial state for the next sequence

Citations

If you use this implementation, please cite FastPLMs and the original ProteinTTT paper:

@misc{FastPLMs,
    author       = {Hallee, Logan and Bichara, David and Gleghorn, Jason P.},
    title        = {FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.},
    year         = {2024},
    url          = {https://huggingface.co/Synthyra/ESMplusplus_small},
    DOI          = {10.57967/hf/3726},
    publisher    = {Hugging Face}
}

@misc{bushuiev2026proteinneed,
    title        = {One protein is all you need},
    author       = {Anton Bushuiev and Roman Bushuiev and Olga Pimenova and Nikola Zadorozhny and Raman Samusevich and Elisabet Manaskova and Rachel Seongeun Kim and Hannes St\"ark and Jiri Sedlar and Martin Steinegger and Tom\'a\v{s} Pluskal and Josef Sivic},
    year         = {2026},
    eprint       = {2411.02109},
    archivePrefix= {arXiv},
    primaryClass = {cs.LG},
    url          = {https://arxiv.org/abs/2411.02109},
}