ProTok: A Unified Protein Tokenizer via Multi-Codebook Quantization

Protein sequences, structures, and functional annotations each capture different aspects of protein biology, but existing tokenizers mostly encode each modality independently, limiting the expressive power of discrete protein representations. A new paper proposes ProTok—a unified protein tokenizer that maps all three modalities into a unified discrete vocabulary through modular multi-codebook quantization.

Why Protein Representations Need Unified Tokenization

In natural language processing, a unified tokenizer is the foundation for building general-purpose language models. Protein research faces a similar problem: sequence data comes with rich textual corpora, structural data relies on 3D coordinates, and functional annotations follow a multi-level label system. The three sources of information are complementary, yet they are typically handled in isolation.

ProTok's core idea is to leverage product quantization to assign an independent codebook group to each modality, then use modality dropout to teach the model to reason under any subset of modalities. This means a single checkpoint can support flexible inference without retraining. It is an important step toward "one model, multiple input types" in protein representation learning.

Product Quantization (PQ) was originally proposed by Jégou et al. in 2011 for vector compression in large-scale image retrieval. The core idea is to split a high-dimensional vector into several low-dimensional sub-vectors, train an independent small codebook for each subspace, and then combine the quantized results of these subspaces to approximate the original vector. Compared with training a single massive codebook, PQ can express exponentially more combinations under the same codebook storage budget, significantly increasing the capacity of discrete representations. ProTok leverages this property to assign separate codebook groups to the sequence, structure, and function modalities, allowing a unified vocabulary to remain compact while accommodating multimodal semantics.

Modality dropout draws on the missing-modality training paradigm in multimodal learning, such as randomly masking one input modality in vision-language models to force the model to learn cross-modal redundant representations. ProTok's innovation lies in combining this strategy with a quantization tokenizer: after randomly dropping one modality's input during training, the shared codebook must still produce meaningful discrete tokens, ensuring inference consistency under any subset of modality inputs. This contrasts sharply with the traditional approach of training separate models for each input combination.

Co-Design of Multi-Codebook Quantization and Modality Dropout

ProTok adopts a modular design: each modality (sequence, structure, function) has its own codebook group, and product quantization decomposes high-dimensional continuous representations into combinations of multiple low-dimensional subspaces, improving representation capacity while maintaining discretization efficiency.

Modality dropout is the key training strategy. During training, the model randomly masks the input of certain modalities, forcing the shared vocabulary to produce meaningful tokens under different modality combinations. This resembles "missing modality" training in multimodal learning, but ProTok combines it with a quantization tokenizer, allowing one model to adapt to different downstream scenarios—structure only, sequence only, or combined inputs.

Performance on Structure Tokenization Benchmarks

Across seven benchmarks, 14 tasks, and 28 test splits in StructTokenBench, ProTok achieves an average rank of 1.4 among pure structure tokenizers, ranks first on 21 splits, and lands in the top two on 27 of 28 splits.

More importantly, these gains do not come at the expense of reconstruction quality. ProTok's tokens reach 0.939 lDDT on the CATH 4.2 dataset, showing that multi-codebook quantization and multimodal training produce tokens that are both reconstructable and informative. This addresses a common concern: whether stronger downstream performance comes at the cost of reconstruction fidelity—ProTok's results show that both can be achieved simultaneously.

lDDT (local Distance Difference Test) is a widely used alignment-free evaluation metric in protein structure prediction and modeling. It measures how well a predicted structure matches a reference structure in terms of local interatomic distances. It computes the fraction of all atom pairs whose distance deviations fall within several thresholds (typically 0.5, 1, 2, and 4 Å) and then averages them, yielding a value between 0 and 1, where higher is better and indicates more accurate local geometry. Unlike RMSD, which relies on global superposition, lDDT is more sensitive to local structural errors and is unaffected by global rotation and translation. Reaching 0.939 lDDT on the CATH 4.2 dataset means that ProTok's discrete tokens can still faithfully reconstruct 3D geometry after decoding, which is key evidence for judging whether a structure tokenizer is "lossy."

StructTokenBench is a benchmark suite for protein structure tokenizers that integrates multiple downstream tasks and data splits for comparing the representation quality of different tokenizers. It covers 7 benchmarks, 14 tasks, and 28 test splits, evaluating dimensions such as structure reconstruction, function prediction, and retrieval. ProTok's average rank of 1.4 and top-two placement on 27/28 splits show that its advantage does not come from an incidental boost on a single task, but rather from stable overall leadership across a multi-task system.

Zero-Shot Functional Retrieval and Generative Capabilities

ProTok's structure tokens comprehensively outperform baseline structure tokenizers on zero-shot functional retrieval, covering CATH, Enzyme Commission (EC), and Gene Ontology (GO) label hierarchies. This means the tokens themselves encode transferable functional semantics rather than merely fitting structural coordinates.

The vocabulary also supports generative tasks. A compact language model trained from scratch jointly generates backbones and sequences on ProTok tokens, with 52% of sampled backbones being designable, compared with only 21% for a parameter-matched GeoBPE baseline. This gap shows that a unified token stream is not only useful for understanding tasks but can also support generative protein design.

CATH, EC, and GO are three complementary label systems in protein functional annotation. CATH classifies protein domains hierarchically (Class, Architecture, Topology, Homologous superfamily), reflecting folding and evolutionary relationships; EC (Enzyme Commission numbers) classifies enzymes into four-digit codes by catalytic reaction type, reflecting biochemical function; and GO (Gene Ontology) uses controlled vocabulary to describe the molecular functions, biological processes, and cellular components of gene products, forming a multi-level directed acyclic graph of labels. Zero-shot functional retrieval means the model has not been trained on the corresponding labels and directly uses the learned token representations to match functional labels—ProTok leads baselines across all three hierarchies, indicating that its tokens carry transferable functional semantics rather than merely compressing structural coordinates.

Designability is a common evaluation criterion in generative protein design: a generated backbone is fed into an independent sequence design model (such as ProteinMPNN), and the predicted sequence is then refolded with a structure prediction tool (such as ESMFold or AlphaFold). If the predicted structure closely matches the original backbone (usually judged by an scRMSD threshold), the backbone is considered "designable." The gap between 52% and 21% indicates that ProTok's unified token stream not only serves understanding tasks but can also directly serve as a representation foundation for generative design.

Implications for Protein Language Models

ProTok's value lies in demonstrating a path: structure, sequence, and function can share the same discrete vocabulary. This opens new possibilities for protein language models—no longer needing separate encoders for each modality, but instead allowing the model to reason over a unified token stream.

Of course, the scale ProTok has validated so far is limited, and the generative tasks also used a relatively small language model. Whether the unified vocabulary can scale to larger parameter counts and more complex multimodal protein tasks remains an open question. Still, its stable performance across multiple benchmarks and its reconstruction metrics already provide a reproducible starting point for future research.

One-line Summary

ProTok uses modular multi-codebook quantization and modality dropout to build a unified protein tokenizer for flexible multimodal inference and generation.

Paragraph Summary

ProTok is a unified protein tokenizer that uses product quantization to assign independent codebook groups to the sequence, structure, and function modalities, and leverages a modality dropout training strategy so that a single checkpoint can support inference over any modality subset without retraining. Across 28 test splits in StructTokenBench, ProTok achieves an average rank of 1.4, ranks first on 21 splits, and places in the top two on 27. Its tokens reach 0.939 lDDT on CATH 4.2, demonstrating that downstream gains do not come at the cost of reconstruction quality. Its structure tokens comprehensively outperform baselines on zero-shot functional retrieval across CATH, EC, and GO, and the unified vocabulary can support generative tasks, with 52% of sampled backbones being designable, far above the 21% GeoBPE baseline. Overall, ProTok provides a reproducible example of protein language models moving toward a "unified token stream, multimodal input" paradigm.