Kanzi

Kanzi is a flow-based protein structure tokenizer developed by Rahul Dilip, Aayush Varshney, Emily Zhang, and David Van Valen at the California Institute of Technology. Structure tokenizers convert three-dimensional protein geometry into discrete or continuous tokens, providing the shared vocabulary that lets multimodal models reason jointly over structure, sequence, and function. Existing tokenizers typically rely on specialized machinery to maintain spatial symmetry — frame-based representations and SE(3)-invariant attention — and this machinery introduces optimization and scalability difficulties.

Kanzi takes a deliberately simpler route. It is a diffusion autoencoder trained with a flow-matching loss, and it makes three substitutions relative to prior tokenizers: it represents structure with global coordinates instead of local frames, it uses a single unified flow-matching objective in place of composite loss functions, and it uses conventional attention instead of SE(3)-invariant attention. The authors show that these choices improve training stability and let a compact model outperform existing tokenizers on reconstruction while using fewer parameters and less compute.

Posted to bioRxiv in October 2025, Kanzi is released as open-source code with a pretrained checkpoint and a from_pretrained loading API, and it demonstrates that symmetry-specialized architectures are not a prerequisite for effective protein structure tokenization.

#Key Features

• Flow-matching diffusion autoencoder: Kanzi is a diffusion autoencoder trained with a single flow-matching loss, replacing the composite loss functions used by earlier tokenizers.
• Global-coordinate representation: The model operates on global coordinates rather than frame-based representations, and uses standard attention instead of SE(3)-invariant attention, simplifying the architecture.
• Parameter efficiency: At roughly 30 million parameters, Kanzi surpasses current tokenizers in reconstruction quality while reducing model size and computational cost.
• Improved training stability: The unified objective and conventional attention improve optimization stability, addressing a practical weakness of symmetry-constrained tokenizers.
• Open checkpoint and simple API: A pretrained checkpoint is provided together with a DAE.from_pretrained() interface for encoding and decoding, under an MIT license.

#Technical Details

Kanzi encodes protein structure into tokens and decodes them back to coordinates using a diffusion autoencoder. Instead of enforcing rotational and translational invariance through frames and SE(3)-equivariant attention, it works directly in global coordinates and trains the decoder as a flow-matching generative process, so a single objective governs reconstruction. The current release provides Cα-only tokenizers and is trained on the standard AFDB-Foldseek clustered dataset restricted to proteins shorter than 256 residues. Despite its small footprint of about 30 million parameters, Kanzi reports better reconstruction than existing tokenizers at lower compute cost. The authors additionally train an autoregressive generative model on Kanzi tokens; this token-based generator outperforms comparable token-based approaches, though it does not match state-of-the-art continuous diffusion models for structure generation. Code is released on GitHub under the MIT license with a downloadable checkpoint; full-backbone tokenizers and generative models are noted as planned extensions, and batch encoding is not yet supported.

#Applications

Kanzi is aimed at researchers building multimodal protein models that need a compact, efficient way to convert structure into tokens for integration with sequence and function. Its from_pretrained API and small size make it practical to drop into token-based pipelines for structure-conditioned generation, structure-aware language modeling, or representation learning without the engineering overhead of symmetry-specialized layers. The accompanying autoregressive generator illustrates how the tokenizer can be used as the front end of a discrete protein structure generation system.

#Impact

Kanzi challenges the assumption that protein structure tokenizers require frame-based representations and SE(3)-invariant attention, showing that a flow-matching diffusion autoencoder over global coordinates can be more stable, smaller, and more accurate for reconstruction. By lowering the parameter and compute cost of tokenization, it makes token-based multimodal protein modeling more accessible. The work is a preprint, and its current scope is bounded: the released tokenizers are Cα-only and trained on sub-256-residue chains, and the associated token-based generator trails continuous diffusion models for de novo structure generation, leaving full-backbone tokenization and stronger generation as stated next steps.