Adaptive Protein Tokenization

This repository contains weights and models for our paper Adaptive Protein Tokenization. We introduce a new form of protein structure tokenization, where each token provides global, higher frequency information about a proteins structure. This departs from the existing paradigm where each token represents a local neighborhood in a protein chain.

APT gone bananas

arXiv • Demo • Blog Post

Quickstart

1. Create a Python environment (3.10+ recommended).
2. Install the package: pip install -e .
3. Verify checkpoint loading:
   • python scripts/verify_checkpoints.py
4. Example usage:

from apt.models import APTLanguageModel, APTTokenizer

tokenizer = APTTokenizer.from_pretrained()
model = APTLanguageModel.from_pretrained()

5. Run the bundled PDB example:
   • python scripts/run_example.py
6. Sample a protein:
   • python scripts/sample_generate.py

Checkpoints

Name	Description	Download link
tokenizer128.pt	Tokenizer weights trained on a maximum of 128 tokens.	Google Drive
lm128.pt	Language model trained on tokens from tokenizer128.pt.	Google Drive
lm128_cond.pt	Language model trained with CATH-A level conditioning (coming soon!).

Missing files are downloaded automatically from the Hugging Face Hub.

Acknowledgments

We use code from several other works: Kanzi, Proteina, and TorchCFM. If you have questions, reach out to rdilip@caltech.edu.

@article{dilip2026adaptive,
  title={Adaptive Protein Tokenization},
  author={Dilip, Rohit and Varshney, Ayush and Van Valen, David},
  journal={arXiv preprint arXiv:2602.06418},
  year={2026}
}