ProteinTTT enables customizing protein language models to one protein at a time for enhanced performance on challenging targets.
To try ESMFold+ProteinTTT in your browser (no installation), see the Hugging Face Space and the Google Colab notebook.
To install locally, please first install the model you are planning to use with ProteinTTT (for example, if you want to use ESMFold+ProteinTTT, install ESMFold according to the official guide first) and then install the proteinttt package itself:
git clone https://github.com/anton-bushuiev/ProteinTTT
cd ProteinTTT
pip install -e .Developer installation and running tests
Install in editable mode with dev dependencies (includes pytest):
git clone https://github.com/anton-bushuiev/ProteinTTT
cd ProteinTTT
pip install -e ".[dev]"Run tests:
pytest -raIn the following example, we use ESMFold+ProteinTTT to predict the structure of a protein. Here, customizing ESMFold with ProteinTTT leads to structure prediction with twice higher pLDDT.
importtorchimportesmimportbiotite.structure.ioasbsiofromproteinttt.models.esmfoldimportESMFoldTTT, DEFAULT_ESMFOLD_TTT_CFG# Set your sequencesequence="GIHLGELGLLPSTVLAIGYFENLVNIICESLNMLPKLEVSGKEYKKFKFTIVIPKDLDANIKKRAKIYFKQKSLIEIEIPTSSRNYPIHIQFDENSTDDILHLYDMPTTIGGIDKAIEMFMRKGHIGKTDQQKLLEERELRNFKTTLENLIATDAFAKEMVEVIIEE"# Load modelmodel=esm.pretrained.esmfold_v1()
model=model.eval().cuda()
defpredict_structure(model, sequence):
withtorch.no_grad():
output=model.infer_pdb(sequence)
withopen("result.pdb", "w") asf:
f.write(output)
struct=bsio.load_structure("result.pdb", extra_fields=["b_factor"])
print('pLDDT:', struct.b_factor.mean())
predict_structure(model, sequence)
# pLDDT: 38.43025# ============ ProteinTTT =============ttt_cfg=DEFAULT_ESMFOLD_TTT_CFGttt_cfg.steps=10# This is how you can modify configmodel=ESMFoldTTT.ttt_from_pretrained(model, ttt_cfg=ttt_cfg, esmfold_config=model.cfg)
model.ttt(sequence)
# =====================================predict_structure(model, sequence)
# pLDDT: 78.69619# Reset model to original state (after this model.ttt can be called again on another protein)# ============== ProteinTTT ===========model.ttt_reset()
# =====================================The package currently implements ProteinTTT for the following protein language models:
| Model name | Reference |
|---|---|
| DPLM2 Bit-based | Hsieh et al., 2025 |
| ProSST | Li et al., 2024 |
| SaProt | Su et al., 2023 |
| ESMFold | Lin et al., 2023 |
| ESM2 | Lin et al., 2023 |
| ProGen2 | Nijkamp et al., 2023 |
| MSA Transformer | Rao et al., 2021 |
See notebooks/demo.ipynb for more usage examples.
The abstract class proteinttt.base.TTTModule implements all core functionality of ProteinTTT, including batch construction, optimization, and logging. To implement ProteinTTT for your custom model, you can subclass TTTModule and implement several methods that specify your model's interface (the alphabet and logic for predicting logits). Optionally, you can also implement logic for evaluating the model during test-time training to select the optimal step. For an example implementation, see how ProteinTTT is implemented for ESMFold in ProteinTTT/proteinttt/models/esmfold.py. The submodule ProteinTTT/proteinttt/models contains other examples.
If you use ProteinTTT in your research, please cite the following paper:
@article{bushuiev2026one,
title={One protein is all you need},
author={Bushuiev, Anton and Bushuiev, Roman and Pimenova, Olga and Zadorozhny, Nikola and Samusevich, Raman and Manaskova, Elisabet and Kim, Rachel Seongeun and St{\"a}rk, Hannes and Sedlar, Jiri and Steinegger, Martin and Tom{\'a}{\v{s}} and Sivic, Josef},
journal={ICLR 2026 (The Fourteenth International Conference on Learning Representations)},
url={https://arxiv.org/abs/2411.02109},
doi={10.48550/arXiv.2411.02109},
year={2026}
}