Flow-Matching-Based Protein Structure Tokenizer Made Efficient and Easy

Zhe Zhang; Yikai Zhang; Jiangtao Feng; Ya-Qin Zhang; Wei-Ying Ma; Hao Zhou

cs.LGSep 27, 2026

Flow-Matching-Based Protein Structure Tokenizer Made Efficient and Easy

Authors: Zhe Zhang, Yikai Zhang, Jiangtao Feng, Ya-Qin Zhang, Wei-Ying Ma, Hao Zhou

Organizations: Institute for AI Industry Research (AIR), Tsinghua University · Department of Computer Science and Technology, Tsinghua University · Fudan University

Abstract

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		$D_{\text{FS}}$
	RMSD $\downarrow$	TM $\uparrow$	RMSD $\downarrow$	TM $\uparrow$	RMSD $\downarrow$	TM $\uparrow$	RMSD $\downarrow$	TM $\uparrow$
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.

AR
	Quality			Diversity
	scRMSD $\downarrow$	scTM $\uparrow$	Des. $\uparrow$	Div. $\downarrow$	Nov. $\downarrow$	$\alpha$ %	$\beta$ %	c%
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 $\uparrow$	PPL/K $\uparrow$	Similarity $\downarrow$
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
$\leq 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	$D_{\text{FS}}$	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 $\downarrow$	AR scRMSD $\downarrow$
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	$D_{\text{FS}}$
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	$D_{\text{FS}}$
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. $\uparrow$	scRMSD $\downarrow$	scTM $\uparrow$	Div. $\downarrow$	Nov. $\downarrow$	$\alpha$ %	$\beta$ %	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. $\uparrow$	scRMSD $\downarrow$	scTM $\uparrow$	Div. $\downarrow$	Nov. $\downarrow$	$\alpha$ %	$\beta$ %	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 (Å) $\downarrow$	TM-score $\uparrow$
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.

Model	Encoding	1-step reconstruction	20-step reconstruction
ProFiT	10.93 GFLOPs	42.07 GFLOPs	0.633 TFLOPs
Kanzi	2.98 GFLOPs	183.5 GFLOPs	2.928 TFLOPs

Appendix

Table 15: Computational cost of Kanzi and ProFiT.

Protein Representation Learning

Protein representation learning builds computational features that capture sequence, structure, function, and molecular relationships.

See all 69 papers

ProTok: Unified Protein Tokenization via Modular Multi-Codebook Quantization
Oct 13, 2026Romain Hardy, Xiaoman Zhang, Pranav RajpurkarProtein Representation Learning Protein Language Models
Department of Biomedical Informatics, Harvard Medical School
Yeti: A compact protein structure tokenizer for reconstruction and multi-modal generation
May 11, 2026Nabin Giri, Steven Farrell, Kristofer E. BouchardProtein Structure Generation Generative Modeling
Lawrence Berkeley National Laboratory
ENSEMBITS: an alphabet of protein conformational ensembles
May 13, 2026Kaiwen Shi, Carlos OliverProtein Language Models Protein Representation Learning
Department of Computer Science, Vanderbilt University · Center for AI in Protein Dynamics, Vanderbilt University · Department of Molecular Physiology and Biophysics, Vanderbilt University

Oct 13, 2026J/K move · Enter open · S save

ProTok: Unified Protein Tokenization via Modular Multi-Codebook Quantization

Romain Hardy, Xiaoman Zhang, Pranav Rajpurkar

Department of Biomedical Informatics, Harvard Medical School

Protein Representation Learning Protein Language Models Product Quantization

Open paper AskPDF arXiv

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.

Abstract

Protein Representation Learning

Explore similar work

ProTok: Unified Protein Tokenization via Modular Multi-Codebook Quantization

Yeti: A compact protein structure tokenizer for reconstruction and multi-modal generation

ENSEMBITS: an alphabet of protein conformational ensembles