ProTok: Unified Protein Tokenization via Modular Multi-Codebook Quantization

Romain Hardy; Xiaoman Zhang; Pranav Rajpurkar

cs.LGOct 13, 2026

ProTok: Unified Protein Tokenization via Modular Multi-Codebook Quantization

Authors: Romain Hardy, Xiaoman Zhang, Pranav Rajpurkar

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

Figure 1: ProTok architecture. Four frozen encoders process global structure, local geometry, sequence, and function inputs. Each encoder output is quantized through a dedicated group of codebooks. During training, modality dropout randomly masks groups and replaces them with learnable mask tokens (illustrated on the function group), forcing the structure decoder and the auxiliary sequence and function decoders to reconstruct from any modality subset. A single-layer fusion module attends across active groups before decoding. At inference, the same checkpoint operates in different modes by selecting which groups are active.

Functional Site Prediction (AUROC %)
Task	Split	ProSST	ESM3	AminoAseed	DPLM-2	ProTok str (ours)
InterPro Bind.	Fold	52.65	44.82	47.29	47.54	49.74
InterPro Bind.	SupFam	91.04	89.33	90.62	90.79	94.53
InterPro Act.	Fold	58.99	59.70	62.13	63.38	64.80
InterPro Act.	SupFam	90.79	88.28	91.86	90.53	96.51
InterPro Cons.	Fold	55.07	55.55	57.09	59.82	58.34

Table 1: Downstream evaluation on StructTokenBench. We evaluate ProTok str against discrete structure tokenizers on functional site prediction (AUROC %), flexibility and dynamics regression (Spearman

\rho

%), zero-shot conformational switch (Spearman

\rho

%), and remote homology (macro

F_{1}

%), via a shared two-layer MLP probe with best-of-12 selection. Bold / underline : best/second-best per row.

	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

Table 2: Zero-shot functional retrieval on CATH 4.2. Top-1 precision (P@1) from mean-pooled structure tokens ranked by cosine similarity, across nine label hierarchies in three families (CATH, Enzyme Commission, Gene Ontology); a query and candidate match when their label sets intersect. Bold / underline : best/second-best per column.

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

Table 3: Ablation study. We ablate ProTok’s encoder groups (top, cumulative) and the orthogonal regularizer (bottom). We report split-averaged TAPE macro

F_{1}

(%), ProteinGLUE AUROC (%), and SimVQ projection effective rank (%). Each variant is trained separately under an identical schedule, so the full row differs slightly from Table 1 .

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 $\pm$ 0.2	85.79	88.37 (DPLM-2)
BioLIP2 Cat.	SupFam	97.8 $\pm$ 0.1	97.70	93.28 (DPLM-2)
ProteinShake	Test	84.3 $\pm$ 0.4	83.77	81.83 (DPLM-2)
ProteinGLUE	Fold	69.5 $\pm$ 0.4	69.75	65.35 (DPLM-2)
ProteinGLUE	SupFam	80.9 $\pm$ 0.3	81.24	74.63 (DPLM-2)

Appendix

Table S1: Seed variability of ProTok str . Mean

\pm

standard deviation over five seeds of the validation-selected probe configuration, with the single-run value reported in Table 1 and the strongest baseline on each split.

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 $\times$ 256	256	64	0.958	0.769	0.939

Appendix

Table S2: Structure tokenizer summary. We compute codebook utilization rate (UR), normalized perplexity (PPL), and native reconstruction lDDT for ProTok str and four structure tokenizers on CATH 4.2.

Codebook group	Modality	Encoder	Vocabulary	UR	PPL	$\cos(\mathbf{z},\mathbf{q})$	Rel. MSE
Global structure	Structure	AminoAseed	4 $\times$ 256	0.926	0.751	0.354	1.291
Local geometry	Structure	ESM-IF1	4 $\times$ 256	0.991	0.787	0.287	1.425
Sequence	Sequence	ESM-2	4 $\times$ 256	0.997	0.940	0.276	1.447
Function	Function	Qwen3	4 $\times$ 256	0.396	0.250	0.236	1.528
Random reference	—	—	—	—	—	0.175	1.650

Appendix

Table S3: Per-group codebook diagnostics. We measure utilization rate (UR), normalized perplexity (PPL), cosine similarity between each per-codebook input

\mathbf{z}

and its codebook lookup

\mathbf{q}

, and the corresponding relative MSE. Values are averaged across each group’s four codebooks. The reference row gives a nearest-of-256 match against random unit vectors.

Figure S1 : Cross-group mutual information. We compute NMI % between the 16 discrete token streams (4 codebooks per group); cells involving the function group use only the 290-chain annotated subset. Darker borders separate the four groups.

	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

Appendix

Table S4 : Quality of sequence-only representations. We compare three sequence-only representations on four StructTokenBench tasks, reporting AUROC % for functional sites and macro

F_{1}

% for remote homology. Bold marks the best per column.

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

Appendix

Table S5: Co-generation on CATH 4.2. Designability (fraction of samples with self-consistency scTM

\geq

0.5 under a redesigned sequence), mean scTM, diversity, and uniqueness for the two ProTok priors and GeoBPE’s generative model at a matched parameter budget. GeoBPE values are as reported in its paper under a different folding oracle; co-designability, which GeoBPE does not report, is given in the text.

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

Appendix

Table S6: Standardized token probes on CATH 4.2. We train probes on each tokenizer’s structure tokens. Inverse folding reports amino acid recovery (AAR); structure prediction reports per-residue lDDT.

		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

Appendix

Table S7 : Codebook group complementarity. We evaluate three nested ProTok inference modes on five StructTokenBench tasks, following the probing setup of Table 1 and varying only the number of input codebooks across modes. We report AUROC % for functional sites and Spearman

\rho

% for ATLAS RMSF. Bold marks the best mode per column.

Protein Representation Learning

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

See all 69 papers

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
Flow-Matching-Based Protein Structure Tokenizer Made Efficient and Easy
Sep 27, 2026Zhe Zhang, Yikai Zhang, Jiangtao Feng +3Protein Representation Learning Discrete Flow Matching
Institute for AI Industry Research (AIR), Tsinghua University · Department of Computer Science and Technology, Tsinghua University · Fudan University
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

May 11, 2026J/K move · Enter open · S save

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

Nabin Giri, Steven Farrell, Kristofer E. Bouchard

Lawrence Berkeley National Laboratory

Protein Structure Generation Generative Modeling Multimodal Generation

Open paper AskPDF arXiv

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.

Abstract

Protein Representation Learning

Explore similar work

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

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

ENSEMBITS: an alphabet of protein conformational ensembles