Flow Autoencoders are Effective Protein Tokenizers
Dilip, Rohit, Zhang, Evan, Varshney, Ayush, Van Valen, David
–arXiv.org Artificial Intelligence
Protein structure tokenizers enable the creation of multimodal models of protein structure, sequence, and function. Current approaches to protein structure tok-enization rely on bespoke components that are invariant to spatial symmetries, but that are challenging to optimize and scale. 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 to-kenizers 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/. The promise of digital biology is to develop machine learning models that are capable of performing a wide range of tasks, from generating novel therapeutics to reasoning about cellular-level processes (Richardson & Richardson, 1989; Cui et al., 2025; Kuhlman & Bradley, 2019). Proteins, which are essential components of almost all biological processes, are a natural target for machine learning approaches to biological perception and generation. A recent exciting advance has been the development of multimodal deep learning models capable of reasoning over protein sequence, structure, and function (Wang et al., 2024; Zhang et al., 2024; Liu et al., 2024; Gaujac et al., 2024; Hayes et al., 2025). These models are enabled by structure tokenizers, which convert the continuous three-dimensional protein structures into a sequence of discrete tokens from a finite vocabulary using vector quantization (V an Den Oord et al., 2017). Training language models on these discrete token sequences unlocks the possibility of truly multimodal biological models that excel across representation and generative tasks. Work done while not at OpenAI. Work done prior to the author's employment at Amazon. These SE(3)-invariant modules, however, can be challenging to both optimize at scale and extend to more diverse biological molecules (e.g., proteins with post-translational modifications, RNA, DNA). Models that can accurately tokenize protein structures without explicitly encoding spatial symmetries may offer improved flexibility and scalability for modeling biology. Such models, however, do not currently exist. Bridging this gap and exploring the performance of non-invariant protein tokenizers is the primary focus of this work.
arXiv.org Artificial Intelligence
Oct-2-2025