Organizations: Institute for AI Industry Research (AIR), Tsinghua University · Department of Computer Science and Technology, Tsinghua University · Fudan University
As the bridge between protein modality and discrete modeling, protein structure tokenization still largely relies on heavily engineered training objectives tailored to specific downstream tasks and large training datasets, which hinders its transfer to broader application scenarios. To address this issue, we propose ProFiT, a lightweight flow matching tokenizer. With simple training strategies that encourage healthy codebook utilization, ProFiT can be trained efficiently and naturally learns semantically meaningful representations without any manual semantic alignment, while achieving reconstruction quality and generalization that match or surpass those of substantially larger tokenizers. We conduct extensive evaluations across a wide range of settings and demonstrate that ProFiT is a plug-and-play tokenizer adaptable to diverse downstream tasks. This study further reveals the significant potential of the flow matching tokenizer paradigm. Our code is publicly available at https://github.com/QDKStorm/ProFiT.
Figures & tables
Figure 1: (Left) Comparison between tokenizer scale and performance. The horizontal and vertical axes denote the overall reconstruction and generation scores (defined in App. A.1 ). The radius of each point is proportional to the logarithmically scaled tokenizer parameter count. Flow-matching–based tokenizers are highlighted in red , while other tokenizers are shown in blue . All figures in this paper follow this visualization convention. (Right) the training framework of ProFiT and its supported downstream applications.
Figure 2: protein reconstruction examples of ProFiT. Gray : ground truth, blue : ProFiT.
Figure 3: Reconstruction results of DPLM2 and ProFiT on T1137s3. Gray : ground truth, purple : ProFiT, blue : DPLM2
CAMEO2022
CASP15
CASP16
DFS
RMSD ↓
TM ↑
RMSD ↓
TM ↑
RMSD ↓
TM ↑
RMSD ↓
TM ↑
IST
3.86
.744
8.64
.688
6.67
.676
3.55
.721
DPLM2
1.87
.903
8.52
.792
8.30
.767
2.07
.872
FoldToken4
1.22
.941
6.47
.845
2.55
.871
1.40
.911
ESM3
1.02
.969
1.99
.918
3.41
.891
0.91
.962
ProTokens
0.94
.963
1.23
.965
1.25
.957
0.89
.958
Table 1: Reconstruction accuracy results. Subscripts denote standard deviations, with values in RMSD columns reported in units of 10−2 and values in TM-score columns reported in units of 10−3 .
Figure 4: ProFiT exhibits stronger generalization, with reconstruction error largely insensitive to both protein length and the OOD level. The OOD level is quantified using pairwise TM-score against the PDB dataset, and is normalized such that 0 corresponds to in-distribution samples and 1 corresponds to the most out-of-distribution samples.
Quality
Diversity
scRMSD ↓
scTM ↑
Des. ↑
Div. ↓
Nov. ↓
α %
β %
c%
AR
DPLM2
8.75 .60
.626 .027
.71 .03
.348 .024
.831 .031
.411 .003
.071 .009
.518 .010
FoldToken4
8.07 .64
.672 .012
.60 .08
.232 .015
.759 .012
.451 .017
.080 .003
.469 .017
ESM3
6.74 .16
.665 .008
.69 .01
.241 .001
.791 .044
.358 .004
.198 .005
.444 .006
IST
6.36 .85
.724 .021
.80 .02
.229 .020
.773 .031
.386 .013
.200 .011
.414 .010
Table 2: Results for unconditional generation with the AR and DPLM2 architectures. Des., Div., and Nov. denote designability, diversity, and novelty, respectively. Standard deviations from five independent runs are indicated in the lower-right corner. Because Kanzi does not provide a backbone-level checkpoint, we cannot evaluate its secondary-structure composition.
Figure 5: Effect of sampling temperature on AR generation quality.
Table 8
Model
UR ↑
PPL/K ↑
Similarity ↓
ProFiT
98.91%
86.83%
0.959
ProTokens
98.24%
66.15%
0.994
ESM3
92.36%
74.70%
0.984
DPLM2
82.76%
63.48%
0.846
Kanzi
71.20%
59.78%
0.986
FoldToken4
52.93%
19.09%
0.987
Table 5: Codebook utilization across all baselines.
Appendix figures & tables13 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Relative performance comparison.
Hyperparameter
Value
Train/validation split
0.999/0.001
Encoder layers
5
Decoder layers
5
Attention heads
8
Token dimension
256
Latent dimension, d
8
Appendix
Table 6: Hyperparameters for ProFiT.
Length regime
CAMEO
CASP15
CASP16
≤256 residues
0.993 / 0.940
0.932 / 0.959
1.341 / 0.939
>256 residues
20.63 / 0.320
36.06 / 0.223
22.91 / 0.343
Appendix
Table 7: Length-stratified performance of Kanzi. Each entry reports RMSD (Å) / TM-score.
Metric
DFS
CAMEO
CASP16
CASP15
Mean TM-score
0.9531
0.8609
0.7559
0.7502
Appendix
Table 8: Mean TM-score of Protenix v1 across evaluation sets.
Method
Recon. RMSD ↓
AR scRMSD ↓
ProFiT
0.84
5.43
- token masking
1.42
5.97
- code replacement
1.45
6.55
Appendix
Table 9: Ablation study on CAMEO2022. We report reconstruction RMSD and AR generation scRMSD.
Figure 7: Ablation studies for Sec. 4.2 .
Figure 8: Performance under different numbers of sampling steps.
Steps
CAMEO
CASP15
CASP16
DFS
2
0.76 / 0.977
4.31 / 0.897
1.78 / 0.953
0.67 / 0.977
5
0.68 / 0.981
2.73 / 0.924
1.34 / 0.968
0.49 / 0.985
20
0.96 / 0.964
1.48 / 0.954
1.41 / 0.952
0.71 / 0.978
Appendix
Table 10: Effect of the number of ODE sampling steps on reconstruction performance. Each entry reports RMSD (Å) / TM-score.
Method
CAMEO
CASP15
CASP16
DFS
Kanzi
9.33 / 0.674
21.86 / 0.520
13.33 / 0.604
0.98 / 0.943
ProFiT ( K=1000 )
0.78 / 0.978
1.38 / 0.964
1.80 / 0.969
0.50 / 0.984
ProFiT ( K=8192 )
0.64 / 0.983
1.19 / 0.967
1.01 / 0.977
0.42 / 0.988
Appendix
Table 11: Codebook-size ablation for reconstruction. Each entry reports RMSD (Å) / TM-score.
Method
Des. ↑
scRMSD ↓
scTM ↑
Div. ↓
Nov. ↓
α %
β %
c%
Kanzi
0.63
12.38
0.429
0.289
0.721
–
–
–
ProFiT ( K=1000 )
0.82
5.80
0.710
0.251
0.712
0.457
0.115
0.428
ProFiT ( K=8192 )
0.82
5.43
0.729
0.262
0.714
0.466
0.106
0.428
Appendix
Table 12: Codebook-size ablation for unconditional generation with the AR architecture.
Method
Des. ↑
scRMSD ↓
scTM ↑
Div. ↓
Nov. ↓
α %
β %
c%
Kanzi
0.48
11.35
0.496
0.346
0.706
–
–
–
ProFiT ( K=1000 )
0.79
6.74
0.660
0.299
0.654
0.621
0.072
0.307
ProFiT ( K=8192 )
0.81
6.25
0.685
0.298
0.655
0.617
0.067
0.316
Appendix
Table 13: Codebook-size ablation for unconditional generation with the DPLM architecture.
Method
RMSD (Å) ↓
TM-score ↑
Kanzi
17.01
0.351
ProFiT ( K=1000 )
8.96
0.750
ProFiT ( K=8192 )
8.31
0.787
Appendix
Table 14: Codebook-size ablation for conditional folding.
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.
Romain Hardy, Xiaoman Zhang, Pranav Rajpurkar
Department of Biomedical Informatics, Harvard Medical School
Multimodal models that jointly reason over protein sequences, structures, and function annotations within a unified representation hold immense potential for integrating multimodal data and generating new proteins with designed functional properties. To utilize transformer architectures, such models require a tokenizer that converts protein structure from continuous atomic coordinates into discrete representations suitable for scalable multimodal training. The quality of such models are fundamentally upper bounded by the fidelity and expressiveness of the underlying tokenized structure. However, existing tokenizers prioritize reconstruction over generative abilities. To address these gaps, we introduce Yeti, a simple and compact protein structure tokenizer based on lookup free quantization and trained end to end with a flow matching objective for multimodal learning. Compared to existing models, Yeti generally achieves the best codebook utilization and token diversity, and second best reconstruction accuracy (with 10x fewer parameters than ESM3) on diverse datasets. To validate Yeti's generative capability, we trained a compact multimodal model jointly over its structure tokens and amino acid sequence entirely from scratch, with no pretrained initialization. The resulting multimodal model generates plausible structures under unconditional cogeneration of protein sequence and structures, achieving comparable results to 10x larger models. Together, these results demonstrate that Yeti is a compact and expressive protein structure tokenizer suitable for training multimodal models that cogenerates highly plausible sequences and structures.
Protein structure tokenizers (PSTs) are workhorses in protein language modeling, function prediction, and evolutionary analysis. However, existing PSTs only capture local geometry of static structures, and miss the correlated motions and alternative conformational states revealed by protein ensembles. Here we introduce Ensembits, the first tokenizer of protein conformational ensembles. Ensembits address challenges inherent to tokenizing dynamics: deriving informative geometric descriptors across conformations, permutation-invariance encoding of variable-size ensembles, and conquering sparsity in dynamics data. Trained with a Residual VQ-VAE using a frame distillation objective on a large molecular dynamics corpus, Ensembits outperforms all related methods on RMSF prediction, and is the strongest standalone structural tokenizer on an token-conditioned ANOVA test on per-residue motion amplitude. Ensembits further matches or exceeds static tokenizers on EC, GO, binding site/affinity prediction, and zero-shot mutation-effect prediction despite using far less pretraining data. Notably, the distillation objective enables Ensembits to predict dynamics token from one single predicted structure, which alleviates dynamics data sparsity. As the field moves from static structure prediction toward ensemble generation, Ensembits offer the discrete vocabulary needed to bring dynamics into protein language modeling and design.
Kaiwen Shi, Carlos Oliver
Department of Computer Science, Vanderbilt University · Center for AI in Protein Dynamics, Vanderbilt University · Department of Molecular Physiology and Biophysics, Vanderbilt University