ProTok: Unified Protein Tokenization via Modular Multi-Codebook Quantization
Organizations: Department of Biomedical Informatics, Harvard Medical School
Abstract
Protein sequence, structure, and functional annotations each capture complementary aspects of protein biology, yet existing tokenizers encode each modality independently, limiting the expressiveness of discrete protein representations. We present ProTok, a unified protein tokenizer that maps all three modalities into a modular discrete vocabulary via product quantization with per-modality codebook groups and modality dropout. A single ProTok checkpoint supports flexible inference from any modality subset without retraining. On seven StructTokenBench benchmarks spanning 14 tasks and 28 test splits, ProTok achieves an average rank of 1.4 across all structure-only tokenizers, placing first on 21 splits and among the top two on 27 of the 28. These gains come without sacrificing reconstruction: ProTok's tokens reach 0.939 lDDT on CATH 4.2, showing that multi-codebook quantization and multi-modal training yield tokens that are both reconstructable and informative. ProTok's structure tokens further lead all baseline structure tokenizers on zero-shot functional retrieval across CATH, Enzyme Commission, and Gene Ontology label hierarchies. The vocabulary also supports generation: a compact language model trained from scratch over ProTok tokens jointly generates backbone and sequence, with 52% of sampled backbones designable versus 21% for a parameter-matched GeoBPE baseline. ProTok shows that structure, sequence, and function can share one discrete vocabulary, motivating protein language models that reason over a unified token stream.
Figures & tables
| Task | Split | ProSST | ESM3 | AminoAseed | DPLM-2 | ProTok str (ours) |
|---|---|---|---|---|---|---|
| Functional Site Prediction (AUROC %) | ||||||
| InterPro Bind. | Fold | 52.65 | 44.82 | 47.29 | 47.54 | 49.74 |
| SupFam | 91.04 | 89.33 | 90.62 | 90.79 | 94.53 | |
| InterPro Act. | Fold | 58.99 | 59.70 | 62.13 | 63.38 | 64.80 |
| SupFam | 90.79 | 88.28 | 91.86 | 90.53 | 96.51 | |
| InterPro Cons. | Fold | 55.07 | 55.55 | 57.09 | 59.82 | 58.34 |
| CATH | EC | GO | |||||||
|---|---|---|---|---|---|---|---|---|---|
| Tokenizer | Class | Arch. | Topo. | EC-1 | EC-2 | EC-3 | MF | BP | CC |
| ProSST | 0.933 | 0.757 | 0.655 | 0.510 | 0.345 | 0.337 | 0.355 | 0.269 | 0.407 |
| ESM3 | 0.924 | 0.740 | 0.634 | 0.498 | 0.390 | 0.361 | 0.359 | 0.244 | 0.350 |
| AminoAseed | 0.888 | 0.701 | 0.632 | 0.506 | 0.373 | 0.321 | 0.379 | 0.269 | 0.350 |
| DPLM-2 | 0.888 | 0.652 | 0.541 | 0.470 | 0.321 | 0.285 | 0.309 | 0.224 | 0.394 |
| ProTok str | 0.945 | 0.868 | 0.817 | 0.695 | 0.562 | 0.502 | 0.490 | 0.407 | 0.516 |
| Variant | TAPE | ProteinGLUE | Rank |
|---|---|---|---|
| Global structure | 64.1 | 68.2 | — |
| Local geom. | 71.1 | 73.6 | — |
| Sequence | 70.4 | 74.9 | — |
| Function (full) | 74.5 | 73.3 | 80.4 |
| Ortho. reg. | 64.1 | 72.6 | 33.3 |
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
| Task | Split | ProTok str (5 seeds) | ProTok str (Table 1) | Best baseline |
| Functional Site Prediction (AUROC %) | ||||
| BioLIP2 Cat. | Fold | 86.8 0.2 | 85.79 | 88.37 (DPLM-2) |
| SupFam | 97.8 0.1 | 97.70 | 93.28 (DPLM-2) | |
| ProteinShake | Test | 84.3 0.4 | 83.77 | 81.83 (DPLM-2) |
| ProteinGLUE | Fold | 69.5 0.4 | 69.75 | 65.35 (DPLM-2) |
| SupFam | 80.9 0.3 | 81.24 | 74.63 (DPLM-2) | |
| Tokenizer | Vocabulary | Dim. | Bits | UR | PPL | lDDT |
|---|---|---|---|---|---|---|
| ProSST | 2,048 | 256 | 11 | 0.997 | 0.871 | — |
| ESM3 | 4,096 | 128 | 12 | 0.946 | 0.517 | 0.986 |
| AminoAseed | 512 | 1,024 | 9 | 0.807 | 0.485 | 0.634 |
| DPLM-2 | 8,192 | 128 | 13 | 0.917 | 0.612 | 0.900 |
| ProTok str | 8 256 | 256 | 64 | 0.958 | 0.769 | 0.939 |
| Codebook group | Modality | Encoder | Vocabulary | UR | PPL | Rel. MSE | |
|---|---|---|---|---|---|---|---|
| Global structure | Structure | AminoAseed | 4 256 | 0.926 | 0.751 | 0.354 | 1.291 |
| Local geometry | Structure | ESM-IF1 | 4 256 | 0.991 | 0.787 | 0.287 | 1.425 |
| Sequence | Sequence | ESM-2 | 4 256 | 0.997 | 0.940 | 0.276 | 1.447 |
| Function | Function | Qwen3 | 4 256 | 0.396 | 0.250 | 0.236 | 1.528 |
| Random reference | — | — | — | — | — | 0.175 | 1.650 |
| BioLIP2 Cat. | InterPro Act. | ProteinShake | TAPE | |||||
|---|---|---|---|---|---|---|---|---|
| Tokenizer | Fold | SupFam | Fold | SupFam | Test | Fold | SupFam | Family |
| Amino acid | 65.23 | 72.56 | 53.64 | 74.68 | 59.72 | 9.39 | 6.27 | 25.64 |
| ESM-2 | 93.11 | 97.76 | 68.66 | 97.30 | 88.31 | 46.14 | 75.98 | 98.72 |
| ProTok seq | 90.56 | 97.66 | 61.02 | 96.99 | 79.79 | 40.46 | 66.31 | 97.16 |
| Prior | Designable % | scTM | Diversity | Unique % |
|---|---|---|---|---|
| GeoBPE | 21.0 | 0.405 | 0.731 | 76.0 |
| ProTok str-seq (75M) | 52.0 | 0.538 | 0.713 | 80.0 |
| ProTok str-seq (156M) | 54.0 | 0.530 | 0.705 | 76.0 |
| Tokenizer | Training | AAR | lDDT |
|---|---|---|---|
| ProSST | Structure-only | 0.143 | 0.636 |
| ESM3 | Structure-only | 0.251 | 0.719 |
| AminoAseed | Structure-only | 0.191 | 0.694 |
| DPLM-2 | Joint | 0.369 | 0.692 |
| ProTok str | Joint | 0.359 | 0.724 |
| BioLIP2 Bind. | BioLIP2 Cat. | InterPro Bind. | InterPro Act. | ATLAS RMSF | |||||||
|---|---|---|---|---|---|---|---|---|---|---|---|
| Mode | No. CB | Fold | SupFam | Fold | SupFam | Fold | SupFam | Fold | SupFam | Fold | SupFam |
| ProTok str | 8 | 78.8 | 83.5 | 85.8 | 97.7 | 49.7 | 94.5 | 64.8 | 96.5 | 52.2 | 47.2 |
| ProTok str-seq | 12 | 83.3 | 87.0 | 91.6 | 98.6 | 51.4 | 96.1 | 65.5 | 96.8 | 53.4 | 49.4 |
| ProTok str-seq-fun | 16 | 83.8 | 88.1 | 93.1 | 98.7 | 56.2 | 97.2 | 67.1 | 98.1 | 48.0 | 39.7 |
Protein Representation Learning
Protein representation learning builds computational features that capture sequence, structure, function, and molecular relationships.
See all 68 papers