Flow Autoencoders are Effective Protein Tokenizers

TL;DR

Kanzi introduces a flow-based protein structure tokenizer that operates directly on backbone coordinates and is trained with a flow matching loss, avoiding SE(3)-invariant components. The architecture uses a lightweight encoder, a diffusion decoder, and a finite quantization bottleneck to produce discrete tokens that condition a diffusion transformer, enabling high-quality reconstruction and scalable autoregressive generation. Across reconstruction and generation benchmarks, Kanzi matches or surpasses larger tokenizers with substantially less data and parameter count, and the autoregressive Kanzi-AR model achieves competitive designability and diversity, underscoring the viability of non-invariant tokenizers for protein modeling. The work also introduces the rFPSD metric for distribution-level reconstruction assessment and discusses limitations and avenues for extending tokenization to larger proteins and full-atom representations, with public code to support reproducibility and broader adoption.

Abstract

Protein structure tokenizers enable the creation of multimodal models of protein structure, sequence, and function. Current approaches to protein structure tokenization rely on bespoke components that are invariant to spatial symmetries, but that are challenging to optimize and scale. We present Kanzi, a flow-based tokenizer for tokenization and generation of protein structures. Kanzi consists of a diffusion autoencoder trained with a flow matching loss. We show that this approach simplifies several aspects of protein structure tokenizers: frame-based representations can be replaced with global coordinates, complex losses are replaced with a single flow matching loss, and SE(3)-invariant attention operations can be replaced with standard attention. We find that these changes stabilize the training of parameter-efficient models that outperform existing tokenizers on reconstruction metrics at a fraction of the model size and training cost. An autoregressive model trained with Kanzi outperforms similar generative models that operate over tokens, although it does not yet match the performance of state-of-the-art continuous diffusion models. Code is available here: https://github.com/rdilip/kanzi/.

Figures (9)

• Figure 1: Schematic overview of our approach. (1) A learned discrete codebook conditions a flow matching model to reconstruct proteins using a diffusion loss. (2) The learned tokens can be used for downstream autoregressive generation, and the generated tokens condition the diffusion decoder to generate protein backbones.
• Figure 2: Architectural overview of Kanzi. (a) Kanzi takes a clean structure as input, which is encoded and passed through a quantization bottleneck. The decoder is provided with the quantized latents as in-context conditioning, along with a noised version of the protein structure. The training is supervised by a single diffusion loss that maximizes $p(\mathbf{x}|\hat{\mathbf{c}})$. No auxiliary losses are used. (b) Our decoder follows the standard diffusion transformer (DiT) presentation, with several notable deviations. We share adaLN conditioning across all blocks, and each DiT block is a transformer with pair-biased attention and optional self-conditioning. (c) Our encoder combines raw coordinate information with sequence positional information. Tokens are mixed using a small stack of transformer layers with sliding window attention. Ablations on other encoder variants are described in Section \ref{['sec:ablations']} and Appendix \ref{['sec:app:ablations']}.
• Figure 3: Designable samples generated from an autoregressive model trained on Kanzi tokens. scRMSDs shown underneath each visualization.
• Figure 4: Left: Scaling of protein structure tokenizer performance with dataset size and parameter count. We plot the reconstruction accuracy on the CAMEO test set versus the training dataset size. Circle area is the model parameter count. Kanzi is competitive with the ESM3 tokenizer, despite a 20-fold smaller parameter count and 400-fold smaller training dataset. Right: Kanzi simplifies the training pipeline, replacing collections of complex, invariant losses with a single, non-invariant flow matching loss.
• Figure 5: All flow losses use the same noise schedule; as numerical values are uninterpretable, we omit them. (a) Encoders need mixing: A point-wise encoder can achieve remarkably strong reconstructions (though underperforms token mixing). Downstream generative performance, however, is poor (see Appendix). (b) Codebook utilization is emergent. Across a large number of architectural classes, codebook usage shoots up after a large number of gradient steps. (c) Invariant representations struggle. Common encoders like MPNNs lead to codebook collapse and identical performance as an unconditional model. While other invariant encoders work, they underperform simply learning the pose. (d) DAEs are scalable. Larger model sizes converge to the same or lower loss in fewer gradient steps. Plot is log-log.