Flow Autoencoders are Effective Protein Tokenizers
Authors: Rohit Dilip
Evan Zhang
Ayush Varshney
David Van Valen
Venue: ICLR 2026 Poster (2026)
Keywords: flow tokenizers
proteins
generation
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: Flow matching tokenizers for better protein structure tokenization and generation

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.

FLOW AUTOENCODERS ARE EFFECTIVE
PROTEIN TOKENIZERS
Rohit Dilip
Evan Zhang
Ayush Varshney
California Institute of Technology rdilipåcal tellip@caltech....
OpenAl
Carnegie Mellon University
evanz@openai.com
varssneeycmu.edd edu
David Van Valen
California Institute of Technology Howard Hughes Medical Institute vanvalen@caltedd edu
ABSTRACT
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 spatial symmetries, but that are challenging to optimize and scale. We present Kanzi, a flow-based to- kenizer for tokenization and generation of protein structures. Kanzi consists of autoencoder trained with flow matching loss. We show that this approach simplifies several aspects of protein structure tokenizers: frame-based representa- tions can be replaced with global coordinates, complex losses are replaced with a single flow matching loss, and SE(3)-invariant attention operations can be re- placed with standard attention. We find that these changes stabilize the training of parameter-efficient models that outperform existing tokenizers on reconstruc- tion metrics at 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 contin- uous diffusion models.
INTRODUCTION
The promise of digital biology is develop machine learning models that are capable of perform- ing a wide range of tasks, from generating novel therapeutics to reasoning about cellular-level pro- cesses (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.] [2244] Liu et al.) 2024 Gaujac et al. 2024] Hayes et al. 2025). These models are enabled by structure tokenizers, which convert the continu- ous three-dimensional protein structures into a sequence of discrete tokens from a finite vocabulary using vector quantization (Van Den Oord et al.) 2017). In these tokenizers, protein structures are represented as tensors R Lxxx3, where L is the sequence length and is the number of backbone atoms. Training language models on these discrete token sequences unlocks the possibility of truly multimodal biological models that excel across representation and generative tasks.
An established practice in protein tokenization is using model components that are invariant to spa- tial symmetries, namely SE(3) (the special Euclidean group, equivalent to SO(3) x Re). This follows the hypothesis that by explicitly encoding inductive biases, models will not hallucinate physically implausible samples that break the symmetry. These SE(33-invariant modules, however, can be chal- lenging to both optimize scale and extend to more diverse biological molecules (c.g., proteins with
"Equal contribution.
Work done while not OpenAT.

---

Learned codebook
volx,,t,
(1) codebook conditions flow model
(2) Autorogressively generate conditioning tokens
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.
post-translational modifications, RNA, DNA). Models that can accurately tokenize protein struc- tures 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. In this work, we describe Kanzi, a flow-based tokenizer for protein structures. Kanzi uses flow matching loss to train autoencoder that tokenizes protein structures. This flow loss simplifies model training by replacing the collection of symmetry-invariant reconstruction losses that commonly used train protein structure tokenizers. Kanzi operates directly the 3D coordinates of backbone atoms and uses standard attention rather than JE(3)-invariant geometric attention methods (Jumper challenge the inductive bias paradigm (Abramson et al. 2024 Geffner et al., (025). Kanzi demon- strates that these simplifications can actually improve tokenizer performance, achieving superior reconstruction quality over larger models trained on more extensive datasets. We next train an au- toregressive model on tokenized structures from Kanzi to sample plausible protein structures. While discrete and continuous diffusion models have seen wide use, autoregressive models are better suited for generating variable-length sequences. This capability critical for tasks like motif scaffolding or in situ structure prediction where the protein sequence (and hence length) is unknown priori. Despite a relatively modest autoregressive model, we match or outperform larger tokenized models in terms of generation quality, measured by standard benchmarks. As supplementary contri- bution, we introduce a reconstruction metric, the reconstruction Fréchet Protein Structure Distance (rFPSD), which utilizes probability divergences to measure structure tokenization. This extends prior work in Faltings et al. (2025) and Geffner et al.] (2025) that applies these metrics to generation. We provide open-source software package as a standalone repository for end-users.
2021 Hayes al. 2025). Consistent with several recent works that
In summary, our primary contributions are as follows:
1. We present Kanzi, a scalable state-of-the-ar flow-based structure tokenizer based on novel asymmetric encoder-decoder design.
2. We demonstrate that simple diffusion loss can replace complex invariant//qquivariant tok- enizer losses, yet still achieve SOTA reconstructions.
3. We Kanzi train an autoregressive protein structure generation model. On standard generative benchmarks, the resulting generations match or outperform existing generative capabilities. To our knowledge, this is the first tokenized model that produces designable structures without massive pretraining.
4. We extend prior work developing distributional metrics for proteins to the reconstruction task to provide broader information tokenization performance. Through series of careful ablations, we demonstrate that while non-invariant encoders can learn scalably, invariant encoders struggle to condition non-invariant decoders.
2

---

(a)
(b)
(c)
Vo(Xe, !)
XN
MLP
Linear
Sliding window attention
Difflssion Loss
DIT Blocks (uN)
MLP
Ember
Sharco
Embbed
Encoder
Liner
Timestep
X1
Seq id
Conditioning
Noized embecking
Figure 2: Architectural overview of Kanzi. (a) Kanzi takes clean structure as input, which is encoded and passed through quantization bottleneck. The decoder is provided with the quantized latents as in-context conditioning, along with noised version of the protein structure. The train- ing is supervised by a single diffusion loss that maximizes p(x(é), No auxiliary losses are used. (b) Our decoder follows the standard diffusion transformer (DiT) presentation, with several notable deviations. We share adaL.N conditioning across all blocks, and each DiT block transformer with pair-biasee bias-d attention and optional ceff-conditioning. (c) Our encoder combines information with sequence positional information. Tokens are mixed using a small stack of trans- former layers with sliding window attention. Ablations other encoder variants described in Section and Appendix
coordinate
2 RELATED WORK
State-of-the-art generative image models frequently first train an image tokenizer, which downsamples continuous image data to either diserete continuous latent (Esser al., 2021 Van Den Oord et al. (2017). Recent works in machine learning for biology have followed suit by training discrete tokenizers for protein backbone structures, which enable language models to be trained on sequences of tokens derived from the resulting codebooks (van Kempen et al., 2022] Steinegger & Söding [017]] Gaujac al.] 20224 Lin al., 2023). While tokenized protein models have underperformed diffusion models on the task of unconditional structure generation, they enable the construction multimodal generative models of proteins. ESM3 notably trained a multimodal discrete diffusion model over sequence, structure, secondary structure, and natural language, which
Tokenization.
was capable generating novel proteins with specified functions (Hayes et al. 2025 Wang et al.) 2024). Following AlphaFold2, protein tokenizers generally rely - -nnvariant architectural components (e.g., invariant point attention) and SE(3) SE(3)-invarian) -invarrat losses (e.g.. frame-ald arame-aligned point error). In contrast with prior work,
non nonninnariant diffusion loss to supervise the tokenizer.
Diffusion and flow matching, State-of-the-an protein structure generation models rely on diffusion, either discrete continuous. FrameDifff, RFDiffusion, and Chroma (Yim al. 2023b Watson ett al. 2023 Ingraham et al. 2023) generate protein backhones using a denoising process over the joint translation-rotation group (3(3) x 3 while FrameFlow. FoldFlow, and FoldFlow-2 similarly perform flow matching over the same manifold (Yim et al.) 2023a Bose et al., [2023 Huguet et al., 2024b). Genie2 and Proteina are more recent attempts to train diflusion models at scale the
AlphaFold Structure Database (AFDB); these both operate over Ca coordinates (Geffner et al.,
2025 Lin al. 2024). The latter does explicitly encode invariances, strategy we broadly adopt here. Despite the strong performance of diffusion models, autoregressive models have unique features that are valuable the structural biology and machine learning communities. Most notably, they can be applied to more
cases where the protein size is not known priori, an important feature for tasks such as motif scaffolding or in situ structure prediction in electron tomography images (Yadav al. 2020 Bunne al. 2024).
To train Kanzi, we use a flow matching objective (Esser et al. 2024 Lipman et al. 2022 Lin
et al. 2024h. Flow matching interpolates between source distribution Po (often Gaussian) and target distribution Paata by integrating along the ODE dx, = (x,,,, t) dt using a learned vector field Nol (xxx))) () IRd (0, 1] R As the vector field generating the true probability distribution is in
3

---

1.27 A
1.86
0.82
2.00
1.99
1.90A
0.91A
Figure 3:
Designable samples generated from an autoregressive model trained on Kanzi tokens.
scRMSDs shown underneath each visualization.
Structure tokenizer model landscape
Dataset sa0 Sagscll
ESM3
dRMSD, direction, binned dRMSD, binned direction
PPLM2
DPLM2
FAPE, violation. histogram
FaldToken
global, fragment, pairwise, pairwise,neighhon neighbor,
blo2token
Kabsch loss, dRMSD
Kant Ooanel
Faidoot
AIDO
FAPE, RMSD, dRMSD
IST
FAPE
bio2tpcen
L75
2.5
Kanzi
Diffusion loss
Figure 4: Left: Scaling of protein structure tokenizer performance with dataset size and parameter count. We plot the reconstruction accuracy the CAMEO test set versus the training dataset size. Circle area the model parameter count. Kanzi is competitive with the ESM3 tokenizer, despite 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.
general unknown, one uses conditional flow matching, which constructs a conditional probability path between prior samples X Po and data samples X1 Pdata- Explicitly, general probability path can be written x, = 0x + d.c, with ~ (0, 1). This induces true conditional vector field u(x, Xg, x1) 6,x1 + o,Xg, where the dot denotes the time derivative. This is a target we can regress against; for the standard case of the linear interpolation path, we have (1 t)xx + fx and u(x; Xo, x1 = x Xg- For completeness, we present a more thorough derivation of the flow matching formulation in Appendix H
Flow autoencoders. The idea of using a flow or diffusion model as the decoder tokenizer re- construction is recent insight in computer vision but has yet to be explored for protein structure generation (Precchakul et al., (022). Two recent works, FlowMo and DiTo (Sargent et al. 2025] Chen al. 2025), both study this approach and independently demonstrate SOTA performance on ImageNet- reconstruction. In these works, the use diffusion model eliminates the need for combinations of perceptual and adversarial losses during training, which were critical insights introduced by VQGAN (Esser et al. 202T). In our case, given the success of recent models like AlphaFold3, Boltz (Wohlwend et al. 2025). and Proteina eschewing symmetric architectures for scalability, we hypothesize that flow autoencoders could provide a similar advantage for tokeniza- tion.
METHOD
We train flow-based tokenizers to perform autoregressive generation. The use diffusion/flow model the decoder allows for considerably more flexibility and scalability in sampling, while the autoregressive prior trained over quantized sequences provides length-agnostic generation.
3.1 ARCHITECTURE
We design a non-equivariant flow autoencoder to tokenize an input structure. It consists of a lightweight encoder ea. substantially deeper decoder dg. and a quantization bottleneck. A pro- tein structure be represented tensor RLXAx3 where is the sequence length and A is
4

---

the number of backbone atoms. Generally A = 1 for Co only models A = 3 for full backbone models. Given x, the encoder processes the raw, mean-centered coordinates using a stack of option- ally pair-biased multi-head self-attention layers. The encoder outputs latent conditioning sequence cx(x) = c E TR Lxd_
For the quantization layer, we adopt finite scalar quantization (FSQ) and discretize to discrete latent è via č = 1e/2 tanh(Linear(())) where & is the number of levels in each dimension of FSQ (Mentzer et al. 2022).. We generally use 8,5,5,5 for an effective codebook size of 1000. We use è to condition the diffusion decoder, which outputs dæ (x,, I, é) (x i, é), with x, the linearly interpolated noise. We pass gradients the encoder using the standard straight- through estimator.
Figure 2 shows all the components of our architecture. First, in contrast with Sargent et al. (2025) which uses separate concatenated modality streams for both the encoder and the decoder, we use single stream for the encoder but concatenated conditioning (i.e., two streams) for the decoder. Because our data is low-dimensional, we found that this was a critical design choice that allowed gradients to efficiently propagate through shallow encoders. In addition, the encoder uses sliding window attention, while the decoder has full bidirectional connectivity. This bias was included to facilitate autoregressive modeling: see Section G for more discussion. Otherwise, the encoder is significantly smaller than the decoder in both width and depth, which is a common design choice for tokenizers. We use relative positional encodings (RoPE) for the decoder, and ablate between absolute and relative positional encodings for the encoder. Finally, in contrast with the standard Diffusion Transformer (Peebles & Xie 2023), we share adaLN weights across layers for time ditioning, which reduces the parameter count by 2 30%. We justify this choice in Appendix G Our full hyperparameter selections are in Appendix F.I
3.2 TRAINING
We optimize the entire tokenizer end-to-end using flow loss
Lnow wtt.ppm...oo00/00.
- (x1 x
è = FSQ (ea(x))
(1)
We again note the relative simplicity of this loss; Figure 4 and Appendix provide detailed description of prior losses supervise tokenization. We train until convergence using AdamW with 81 0.9, 32 0.95, and learning rate / = 1.7 10-2 We use linear warmup and cosine decay schedule with random rotations on the inputs as augmentations. Details of our train- ing/hyperparameter configurations are given in Appendix [F.1] One advantage of flow autoencoders is the ability to use classifier-fret free guidance to improve sample quality, We mask out the condition- ing sequence è with probability 0.1 to enable this option. Following AlphaFold3 and Proteina, mean-center all proteins using Ca coordinates and augment input structures with random rotations during training.
3.3 DATASET
While early protein generative models trained primarily the ~ 30k distinct structural homologs in the Protein Data Bank (Bank) 197T), recent works like Proteina and AlphaFold3 have trained on synthetic AlphaFold2 predictions to achieve significant performance gains. As Proteina and AlphaFold3 were trained on datasets that require extensive computational resources (Proteina trains on the full 214M structures in the AFDB, and AlphaFold3 trains on 40M MGnify structures), instead train on the Foldseek clustered AlphaFold database, which we denote as Drs. Drs filters and clusters the AFDB using MMseqs2 and Foldseek and keeps a single representative structure per cluster. Both Proteina and Genie2 include Dps in their training data. We perform additional filtering (described in Appendix F.2 which leaves total of 498,900 structures.
3.4 INFERENCE
Diffusion decoding. Given a sequence of Kanzitokens, to sample the full distribution we can simply run Euler inference using the standard integrator, i.c., we repeat the following for fixed and e [0, 1/N, 2/N, (N N)/N..
X+ = x 1,(
(2)
5

---

where è is conditioning sequence. We make several noteworthy additions. A major strength of flow tokenizers their ability to take advantage of sophisticated sampling strategies at inference time compared to autoregressive discrete diffusion samplers. We construct new flows using classifier-free guidance with guidance parameter as follows:
(x, 1, c, = vo(x,, (, + g(vo (, vo(x,,t, 2))
We omit the tilde for the remainder of the paper, with the understanding that tune classifier-free guidance unless explicitly stated otherwise. As we use a Gaussian flow, we also have a closed-form expression for the corresponding score field
Sp Xt, t, é) tVo, t. X,
(4)
1-t
An analogous expression for the score field with classifier-free guidance applies. A frequent practice in biological diffusion is to construct an ad hoc stochastic differential equation as an alternative sampler, in Equation
dx. Vo (xt, t, dt ggldttggtnsseeeeeete (( t, dt + /2g(t)
(5)
This notation largely follows that Getfner et al. 2025). Setting y = 1 corresponds to standard Langevin dynamics. However, treating the noise scale y and the score scale 17 as hyperparameters can significantly improve generation quality, although sometimes at the cost of diversity. In our benchmarks, we provide sampling using both the full distribution (e.g., = 7 = 0) and using noise and score-scaling as references.
Autoregression. We explored both nucleus sampling and min-p sampling for generating the con- ditioning sequence c, but did not observe any significant difference between the two. All presented results use nucleus sampling with a cutoff of 0.9, As our generative models are autoregressive, we can also do best-of-N sampling with the log-likelihood as an inexpensive proxy for decoding quality, a strategy that has proven useful for large language models (Qiu et al. 2024: Song et al. 2024 Generally N = 2 or 4 in our experiments.
4 EXPERIMENTS
We evaluate Kanzi on reconstruction, generative, and representative tasks. On reconstruction tasks, we find that flow autoencoders exceed or match the performance much larger models trained on substantially more data. On generative tasks, we exceed match the performance of larger tok- enized models and consistently outperform comparably sized generative models trained on other tokenizers. We defer evaluations representation quality to Appendix D in summary, Kanzi out- performs similar non- con-invariant tokenizers on residue-level representation tasks but underperforms invariant tokenizers.
4.1 FLOW TOKENIZERS ARE SOTA AT RECONSTRUCTION
We first evaluate Kanzi on reconstruction. We benchmark reconstruction performance against all structure tokenizers with accessible public repositories: ESM-3, DPLM-2, Bio2Token, FoldToken, and the InstaDeep Structure Tokenizer (IST). A challenge with evaluating tokenizers is every model tends to use a different training/test set, which makes it challenging compare model performance. To address this, we use wide range of test datasets (all are held-out from our model, along with any structural homologs at 80% similarity as determined by Foldseek). We exclude any cases where know there is leakage between the benchmark set and the tokenizer. See Appendix 1.2 for more details on these determinations. We use five held-out test datasets: CAMEO, CASP 14, CASPI5, CATH, and held-out subset of Dys- We use RMSD and TM-score benchmark local and global measures of reconstruction, and include scores for both full backbone tokenizers and Ca only tok- enizers. We also introduce the following two auxiliary metries.
rFPSD (reconstruction Fréchet Protein Structure Distance) provides a distribution- metric for reconstruction by using the deep features in pretrained CATH-classifier Metries like RMSD

---

CAMEO
CASP14
CASPES
CATH
AFDB
RMSD (4)
TM (t)
RMSD
TM
RMSD
TM
RMSDD
TM
[
[e|RMSD
(FPSD
RMSD
TM
DPLM2 (118M)
1.651
0.876
1,008
0.951
2.160
0.866
1.641
0.897
1.851
2.067
5,742
4.676
0.810
ESM3 (648M)
0.860
0.955
0.462
0.987
1,021
0.969
0.957
1.391
23.399
2.384
0.915
FoldToken (85M)
2,539
0.881
2.194
0.936
6.629
0.744
1.298
0.920
1.575
1.231
71,786
2.161
0.858
IST (1IM)
1.637
0.916
0.900
0.960
1.252
0.953
1.201
0.940
1.127
1.246
105.208
2.872
0.862
bio2token (LIMM)
1.076
0.948
1,006
0.952
1.377
0.939
1.361
1.212
0.932
Kanzi (30M)
0.936
0.948
0.861
0.958
1.345
0.951
1.098
0.940
0.774
1.181
27.202
1.069
0.947
Kanzi (30M)*
0.817
0.960
0.698
0.972
1.267
0.963
0.953
0.955
0.658
L.023
7.956
0.870
0.962
Kanzi (1IM)
1,016
0.937
0.912
0.954
1.259
0.955
156
0.934
0.805
1.239
20.104 104
1.210
0.934
Kanzi (1IMY*
0.863
0.952
0.762
0.968
1.105
0.965
0.994
0.950
0.813
1.058
51.649
0.994
0.952
Table 1: Reconstruction metrics across tokenizers for Ca reconstruction. Kanzi consistently matches or outperforms much larger models trained on larger datasets. Best result in bold, second best result underfined. Datasets like CAMEO and CASP are relatively small and have larger vari- ances. We exclude any cases where a model is explicitly stated to be trained on a held-out dataset. Starred (*) Kanzi models use 17 0.45, 7 1.0,9 g 2.0. This parameter setting is deliberately underoptimized; we tuned against a small subset of our AFDB test set (100 structures) and applied it without adjustment to the held-out non-synthetic datasets. For visual clarity, standard errors are left to the Appendix.
CAMPO
CASP14
CASPIS
CATH
AFDR
RMSD(J) TM (t)
RMSD
TM
RMSD
TM
RMSD
TM
RMSD
TM
DPLM2
1631
0.928
0.959
2.144
1.717
0.925
4.646
0,880
ESM3
0,861
0.980
0.463
0.994
1.018
0,983
1.151
0.971
2,378
0.944
FoldToken
2.498
0.922
2.323
0.958
6,580
0.809
1.352
0,944
2,120
0.907
1.626
0.954
0.896
0.974
1,244
1.195
0.956
2,859
0.904
bio2token
1.069
0.963
0.998
0.966
1.367
0,960
0.987
0.958
1,201
0.949
Kanzi (30M)
0.996
0,973
0.889
0
0,980
1.074
0.972
1.165
0.969
Table 2: Full backbone reconstruction. While the gap is less pronounced than Cox only, Kanzi con- sistently achieves the best or second-bess reconstruction across datasets. rFPSD is a Co only metric by construction, so it is excluded.
biased towards the capabilities of current folding models, and recent work has shown that RMSD is not strong predictor of generative capabilities. rFPSD captures the statistics of the distribution, it may be a useful addition to the metrics used for tokenizer devel- opment. rFPSD extends the original FPSD metric introduced for measuring generative capabilities in (Geffner et al. 2025).
and our results in
[ss]RMSD considers RMSD only over proteins with > 60% of particular secondary structure (ss=a, 3, c) content. This isolates the failure modes of tokenizers and can help measure a tok- enizer's ability model unstructured regions, which is particularly important for many downstream therapeutic tasks.
We consider auxiliary metrics only over the CATH dataset, for two reasons. We do not want to use the AFDB, as we do not want to further bias metrics towards distributions over synthetic struc- tures. Second, most other natural protein structure datasets are not large enough to provide accurate estimators for distribution metrics such as rFPSD.
We train both Ca only and full backbone tokenizers, as both are useful depending on the task at hand. These results are shown in Tables [ and [ Kanzi consistently performs best or second best across dataset categories, despite being trained entirely on synthetic data. The gap more modest for full-backbone tokenization, but we again note the significant difference in model and data scale between Kanzi and structure tokenizers like ESM3 and DPLM2 (see Fig The gap in rFPSD between ESM3 and DPLM2 underscores the value of distribution-leve metrics for assessing and improving tokenizer performance: while DPLM2 underperforms ESM3 on reconstruction, surpasses it rFPSD. As prior work has emphasized that strong reconstruction does not necessarily imply high-quality generation (Geffner et al. 2025 Hsich et al., 2025). developing better metrics to capture generative ability during tokenizer training an important contribution.
7

---

4.2 GENERATION
The primary purpose of tokenization is for downstream representations. To evaluate Kanzi's per- formance on structure generation, we train an autoregressive model on tokenized sequences from Kanzi. To broaden the scope of our comparisons, the DPLM and ESM tokenizers. The full details on these models is in Appendix 1.2 These compar- isons ensure that strong results from model (e.g., autoregressive vs discrete diffusion). The full results for ESM3, DPLM-2, ESM3-AR, and DPLM2-AR are presented in Table El Despite having significantly smaller parameter count than alternative models, Kanzi-AR exhibits strong performance across quality and diversity metrics.
also train additional autoregressive models on
tokenizer can be decoupled from the choice of generative
To demonstrate unique value of the non-invariant approach, train generative model Kanzi tokens conditioned on a cryoET volume map. This is an instance where other tokenizers immediately collapse. as the conditioning signal is intrinsically non-invariant. Appendix D describes the task, demonstrates our approach, and provides quantative evidence of the utility of Kanzi tokens.
Kanzi-AR overpredicts alpha helices, known issue for models that heavily rely on synthetic data. This can be resolved by additional post-training steps (as done in Huguet et al. (2024a)), which we leave future work. We also noticed that evaluations across models have a substantial amount of variance depending on sampling; for instance, our benchmarking on DPLM-2 outperforms the scRMSD results reported in Wang et al. (2024) but underperforms the reported scTM values. To ensure reproducibility, we describe our generation and benchmarking process for all models in Ap- pendix L.2
Model
QUALITY
DIVERSITY
Designahility (f)
scMMDD (J)
scTM (1)
Diversity (1)
Navelty (1)
a%
.%%
c%
ESM3 (650M/32)
0.476
10,898
0.702
0.705
0.730
82.9
5.1
11.9
ESM3 (650M/256)
0.460
6.959
0.7199
0.604
0.692
74.9
11.5
13.6
ESM3-AR (30M)
0.520
4.252
0.804
0.241
0.751
38.6
16.8
44.6
DPLM2 (650M)
0.486
3.314
0.814
0.263
0.735
42.5
15.2
42.3
DPLM2-AB AR (300M)
0.320
8.989
0.706
0.308
0.772
41.2
18.0
40.8
Kanzi-AR (250M)
(1 00
0.328
4.210
0.724
0.271
0.715
71.9
6.5
21.6
(1 0.66)
0.562
3.781
0.795
0.408
0.773
88.7
0.7
10.7
(1 0.66, BoN)
0.617
3.655
0.807
0.386
0.763
88.2
0.8
11.0
Table 3: Generative evaluation across models. Bold best (per column), underline second-best The first ESM3 records are with 32 and 256 steps, respectively. (i.e., for the latter, every pass de- codes single new token). Kanzi-AR consistently shows strong performance an additional contribution, we use best-of-N sampling (N = 2) with log-likeihhoods as our reward proxy to improve performance, demonstrating that our autoregressive prior learns meaningful dis- tribution.
metrics. As
4.3
ABLATIONS AND DESIGN CHOICES
This section contains a curated list of ablations and design choices we made. We present a more thorough list of findings in the appendix. Each bolded item has corresponding reference in Fig- ure 5
Encoders need token mixing for generation, but not for reconstruction. When operating on coordinates without any invariant transforms, because the input itself carries raw positional infor- mation, it isn't obvious that pooling of local information should occur in the same way as it does in image tokenizers or invariant protein structure tokenizers (see Section K.2 for more discussion and Ellmen et al. (025)) for deeper investigation). We ablate the window size encoder trans- formers full attention, window size 8, window size 4, and window size (corresponding to a point-wise MLP on the coordinates). Surprisingly, the latter suffices for good reconstructions, but substantially harms downstream generations.
Codebook utilization is emergent. A surprising but reproducible phenomenon the emergence of high codebook utilization after extended training. At the start of training, the raw coordinates highly correlated, which in FSQ leads low usage. Over time, this spreads out, signifying
8

---

(a)
[b)
(c)
(Large) WS-4 (Mteaimm) WS-4 (Smail)
Windowe8
Codsbook uttiiztton
M&C
Nonw
Window-4
Flow loss
Linear
Nlw loss
100
Flow Iosa
Uioo looa
-1A
0.5-
10-1-
1.1
102
10f
10
0
50000
100000
20000
40000
304
100
10P
Gradient steps
Gradient
Gradent steps (250k)
stepe
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 recon- structions (though underperforms token mixing). Downstream generative performance, however, is poor (see Appendix). (b) Codebook utilization is emergent. Across large number of archi- tectural 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) Flow autoencoders are scalable. Larger model sizes converge the same or lower loss in fewer gradient steps. Plot is log-log.
an increase in codebook utilization (see Appendix G for additional approaches we tried to increase codebook usage).
Invariant representations struggle. Most existing tokenizers use an invariant encoder; we exper- iment with two variants of this design choice that is, we use an invariant encoder to inform an equivariant decoder. We experiment with a learned relative rotation between the encoder and de- coder, which encourages the model to learn an invariant representation, and an explicitly invariant input. Both cases underperform simply allowing the model to tokenize the pose for both recon- struction and generation. Most graph-based invariant models (like MPNN, a common choice for proteins) lead to codebook collapse and the same flow loss as a purely unconditional model.
Flow autoencoders are scalable tokenizers. An advantage of transformer layers that use standard attention is their scalability. While this is computationally challenging to explore, small-scale ex- periments we performed (<0.2B parameters) consistently showed that larger models reach the same performance as smaller models after fewer steps, and model performance continues improve with extended training. We observed no evidence of overfitting, and continued to see performance im- provements well over 100k gradient steps during training.
5 OUTLOOK AND LIMITATIONS
In this work, we present a new approach for structure tokenization. We demonstrate that for the task of protein structure tokenization, flow autoencoders that utilize standard attention simplify model training while enjoying performance equal to or better than other state-of-the-art generative models. An autoregressive model trained on Kanzi tokens to our knowledge, the first tokenized structure model competitive with structure tokenizers like ESM3 or DPLM2 that use large-scale pre-training. Additional work necessary to close the performance gap between models that leverage tokeniza- tion for structure generation, as described here, and state-of-the- th-tt--atttttttt art diffusion models.
Our work has several key limitations. We primarily train on the AFDB, which synthetically biased; our introduction of the rFPSD metric was an attempt to measure these effects. For computational reasons, our models are quite small. It challenging to know how these models will actually per- form as data size, model size, and training time are increased, as was studied in|Geffner et al. (2025). Similarly, we only train on proteins of size < 256. Geffner et al. (2025) showed that a fine-tuning stage with larger proteins could dramatically improve long protein designability. We expect that with appropriate computational resources, our models might realize similar gains. Our generative models are trained on Ca tokenizers, choice driven by computational limitations. During train- ing, we were surprised to observe only small differences between Ca and full backbone tokenizers. Extending generation to the full-backbone and all-atom case is important future direction.

---

We close by noting that while diffusion/flow matching models remain state-of-the-arr for protein generation, tokenization has much to offer structural biology and protein design. Biology is filled with diverse representations cryoEM and cryoET images, single cell data, multiplexed immunoflu- orescence data, structural and natural language descriptions of proteins, etc. Tokenized representa- tions are uniquely amenable multimodal tasks beyond generation. We believe that much like the interplay between images, text, video, and audio, building large foundation models for the life sciences will require robust tokenized representations across data modalities.
ACKNOWLEDGMENTS
R.D. thanks Samuel Arnesen, Markus Marks, Laura Luebbert, Sinan Ozbay, William Ameeenn Nicholas Ritter. and Gohta Aihara for technical support and helpful feedback on the draft. R.D. thanks Atlantic Labs for insightful conversations. R.D. thanks Jeff Song and Zhen Chen for discus- sions on the cryoFT methods described in this paper. We thank Modal for supporting this work via a Modal compute grant. This work was supported by an in-kind Landing Fellow Ship.
REPRODUCIBILITY STATEMENT
have taken several steps to ensure reproducibility of our results. Our training relies entirely on public data (Varadi et al.) 2022). The main text describes the core model architecture (Section 3.1 and training objective. We fully describe all hyperparameters in Appendix[F]] F.I and dataset processing steps in Appendix [22]] We describe
evaluation metrics with code references in Appendix C.I
and Appendix C.2 Source code and instructions for reproducing all experiments will be released publicly after the review period, anonymization and cleanup complete.
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630(8016):433500,, 2024.
Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic inter- polants. arXiv preprint arXiv:2209.155711 2022.
Protein Data Bank. Protein data bank. Nature New Biol, 233(223):10-1038 1971.
Avishek Joey Bose, Tara Akhound Akhound-Sadegh, Sadegh, Guillaume Huguet, Kilian Fatras. Jarrid Rector-Brodd Cheng-Hao Liu, Andrei Cristian Nica, Maksym Korablyov. Michael Bronstein, and Alexan- der Tong. Se (3)-stochastic flow matching for protein backbone generation. arXiv preprint arXiv:2310.023991 2023.
Charlotte Bunne, Yusuf Roohani, Yanay Rosen, Ankit Gupta, Xikun Zhang, Marcel Roed, Theo Alexandrov, Mohammed AlQuraishi, Patricia Brennan, Daniel B Burkhardt, al. How build the virtual cell with artificial intelligence: Priorities and opportunities. Cell, 187(25):7045-7663 2024.
Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu, al. Pixart: Fast training of diffusion transformer for photoreal- istic text-to-image synthesis. arXiv preprint 2rX.v:2310.00426. 2023.
Shoufa Chen, Mengmeng Xu, Jiawei Ren, Yuren Cong, Sen He, Yanping Xie, Animesh Sinha, Ping Luo, Tao Xiang, and Juan-Manuel Perez-Rua. Gentron: Diffusion transformers for image and video generation. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 6441-6451, 2024.
Ting Chen, Ruixiang Zhang, and Geoffrey Hinton. Analog bits: Generating discrete data using diffusion models with seff-conditioning. arXiv preprint arXiv:2208. 02402, 2022.
Yinbo Chen, Rohit Girdhar, Xiaolong Wang, Sai Saketh Rambhatla, and Ishan Misra. Diffusion autoencoders scalable image tokenizers. arXiv preprint arXiv:2501. 18593, 2025.
10

---

Haotian Cui, Alejandro Lapuaa-apperta, Maria Brbić, Julio Saez- Simona Cristea, Hani Goodarzi, Mohammad Lotfollahi, Fabian J Theis, and Bo Wang. Towards multimodal foundation models in molecular cell biology. Nature, 640(8059):623 633, 2025.
Justas Dauparas, Ivan Anishchenko, Nathamiel Bennett, Hua Bai, Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob de Haas, Neville Bethel. et al. Robust deep learning- based protein sequence design using proteinmpnn. Science, 378(6615):49)
2022.
Isaac Ellmen, Constantin Schneider, Matthew IJ Raybould, and Charlotte M Deane. Transformers trained on proteins can learn to attend to euclidean distance. arXiv preprint arXiv:2502.01333, 2025.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recogni- tion, pp. 11873-12883 2021.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first international conference machine learning, 2024.
Felix Faltings, Hannes Stark, Tommi Jaakkola, and Regina Barzilay. Protein fid: Improved evalua- tion of protein structure generative models. arXiv preprint WrXi:22505.08041 2025.
Christopher Fifty, Ronald Junkins, Dennis Duan, Aniketh Iyengar, Jerry W Liu, Ehsan Amid, Sebastian Thrun, and Christopher Ré. Restructuring vector quantization with the rotation trick. arXiv preprint 4rXiv:2410.06424, 2024.
Ruiqi Gao, Emiel Hoogeboom, Jonathan Heek, Valentin De Bortoli, Kevin P. Murphy, and Tim Salimans. Diffusion meets flow matching: Two sides of the same coin, 2024. URL https:
/ddiffusionflow. www.ssionnnoo.ggtuu...o.oo io/ Blog post.
Zhangyang Gao, Cheng Tan, Jue Wang, Yufei Huang, Lirong Wu, and Stan Zi. Foldtoken: Learn- ing protein language via vector quantization and beyond. In Proceedings of the AAAI Conference Artificial Intelligence, volume 39, pp. 219-227. 2025.
Victor Garcia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. E (n) equivariant normalizing flows. Advances in Neural Information Processing Systems, 34:4181- 4192, 2021.
Benoit Gaujae, Jérémie Donå, Liviu Copoiu, Timothy Atkinson, Thomas Pierrot, and Thomas Barrett. Learning the language of protein structure. arXiv preprint arXiv:2405.15840. 2024.
Tomas Geffner, Kieran Didi, Zuobai Zhang, Danny Reidenbach, Zhonglin Cao, Jason Yim, Mario Geiger, Christian Dallago, Emine Kucukbenli, Arash Vahdat, al. Proteina: Scaling flow-based protein structure generative models. arXiv preprint 2rXiv:2503.00710 2025.
Thomas Hayes, Roshan Rao, Halil Akin, Nicholas J Sofroniew, Deniz Oktay, Zeming Lin, Robert Verkuil, Vincent Q Tran, Jonathan Deaton, Marius Wiggert, et al. Simulating 500 million years of evolution with a language model. Science, 3876673)):850-888 2025.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:684006851, 2020.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffu- sion for molecule generation in 3d. In International conference machine learning, pp. 8867- 8887. PMLR, 2022.
Cheng- Yen Hsich, Xinyou Wang, Daiheng Zhang, Dongyu Xue, Fei Ye. Shujian Huang, Zaixiang Zheng, and Quanquan Gu. Elucidating the design space of multimodal protein language models. arXiv preprint ar 2504 1454, 2025.
11

---

Chloe Hsu, Robert Verkuil, Jason Liu, Zeming Lin, Brian Hie, Tom Sercu, Adam Lerer, and Alexan- der Rives. Learning inverse folding from millions of predicted structures. In International con- ference machine learning, pp. 8946-8970 PMLR, 2022.
Guillaume Huguet, James Vuckovic, Kilian Fatras, Eric Thibodeau-Laufer, Pablo Lemos, Riashat Islam, Chenghao Liu, Jarrid Rector-Brooks, Tara Akhound-Sadepa Michael Bronstein, et al. Sequence-augmennte augmented (3)-flow matching for conditional protein generation. Advances in neural information processing systems, 37:33007 77330007-33036, 2024a
Guillaume Huguet, James Vuckovic, Kilian Fatras, Eric Thibodeau-Laun Laufer, Pablo Lemos, Riashat Islam, Chenghao Liu, Jarrid Rector-Brooks Tara Akhound-Sadepa Michael Bronstein, et al. Scquencc-auggmented se (3)-flow matching for conditional protein generation. Advances in neural information processing systems, 37:33007-33036 2024b.
John B Ingraham, Max Baranov, Zak Costello, Karl W Barber, Wujie Wang, Ahmed Ismail, Vincent Frappier, Dana M Lord, Christopher Ng- -g-Thow-Hinn, Erik R Van Vlack, et al. Illmiinating protein space with programmable generative model. Nature, 623(7989):1070-1077)
2023.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool. Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583-58) 2021.
Brian Kuhlman and Philip Bradley. Advances in protein structure prediction and design. Nature reviews molecular cell biology, 20(1 11:681-697, 2019.
Xiaohan Lin, Zhenyu Chen, Yanheng Li, Xingyu Lu, Chuanliu Fan, Ziqiang Cao, Shihao Feng, Yi Qin Gao, and Jun Zhang. Protokens: A machine-learned language for compact and informative encoding of protein 3d structures. 2023.
Yeqing Lin, Minji Lee, Zhao Zhang, and Mohammed AlQuraishi. Out of many, one: Design- ing and scaffolding proteins the scale of the structural universe with genie 2. arXiv preprint arXiv:2405. 15489, 2024.
Yaron Lipman, Ricky TQ Chen, Heli Ben- Ren-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Andrew Liu, Axel Elaldi, Nathan Russell, and Olivia Viessmann. Bio2token: All-atom tokenization of any biomolecular structure with mamba. arXiv preprint arXiv:2410.19111. 2024.
Fabian Mentzer, David Minnen, Eirikur Agustsson, and Michael Tschannen. Finite scalar quantiza- tion: Vq-vae made simple. arXiv preprint 2.3iv:2309.15500. 2023.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEF/CVF international conference on computer vision, pp. 4195 4195-4205, 2023.
Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwongsa, and Supasorn Suwajanakorn. Dif- fusion autoencoders: Toward a meaningful and decodable representation. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 10619-10629, 2022.
Jiahao Qiu, Yifu Lu, Yifan Zeng, Jiacheng Guo, Jiayi Geng, Chenhao Zhu, Xinzhe Juan, Ling Yang, Huazheng Wang, Kaixuan Huang, et al. Treebon: Enhancing inference-time alignment with speculative tree-search and best-of-n sampling. arXiv preprint arXiv:2410.1600 2024.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language under- standing by generative pre-training. 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 1(8):9, 2019.
Janes S Richardson and David C Richardson. The de novo design of protein structures. Trends in biochemical sciences, 14(7):304-309 1989.
Kyle Sargent, Kyle Hsu, Justin Johnson, Li Fei-Fei, and Jiajun Wu. Flow the mode: Mode-seeking diffusion autoencoders for state-of-the-ae image tokenization. arXiv preprint arXiv:2503. 2025.
12

---

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456. 2020.
Yifan Song, Guoyin Wang, Sujian Li, and Bill Yuchen Lin. The good, the bad, and the greedy: Evaluation of Ilms should not ignore non-determinism. arXiv preprint arXiv:2407.10457, 2024.
Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Harmonic self-conditioned flow matching for multi-ligand docking and binding site design. arXiv preprint arXiv:2310.05764, 2023.
Martin Steinegger and Johannes Söding. Mmseqs2 enables sensitive protein sequence searching for the analysis of massive data sets. Nature biotechnology, 55(11))0026-0088 2017.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Michel van Kempen, Stephanie S Kim, Charlotte Tumescheit, Milot Mirdita, Cameron LM Gilchrist, Johannes Süding, and Martin Steinegger. Foldseek: fast and accurate protein structure search. Biorxiv, pp. 2022-02, 2022.
Mihaly Varadi, Stephen Anyango, Mandar Deshpande, Sreenath Nair, Cindy Natassia, Galabina Yordanova, David Yuan, Oana Stroe, Gemma Wood, Agata Laydon, et al. Alphafold protein structure database: massively expanding the structural coverage of protein-sequence space with high-accuracy models. Nucleic acids research, 503D1):D499-D44, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
Xinyou Wang. Zaixiang Zheng, Fei Ye, Dongyu Xue, Shujian Huang, and Quanquan Gu. Dplm-2: A multimodal diffusion protein language model. arXiv preprint 2XXx:22410.13822 2024.
Joseph L Watson, David Juergens, Nathaniel Bennett, Brian L. Trippe, Jason Yim, Helen E Eise- nach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature, 62097976):089-11000 2023.
Jeremy Wohlwend, Gabriele Corso, Saro Passaro, Noah Getz, Mateo Reveiz, Ken Leidal, Wojtek Swiderski, Liam Atkinson, Tally Portnoi, Itamar Chinn, et al. Boltz-1 democratizing biomolecular interaction modeling. BioRxiv, pp. 2024-11, 2025.
Pinku Yadav, Olivier Rigo, Corinne Arvieu, Emilie Guen, and Eric Lacoste. In situ monitoring systems of the slm process: On the need to develop machine learning models for data processing. Crystals, 10(6):524, 2020.
Jason Yim, Andrew Campbell, Andrew YK Foong, Michael Gastegger, José Jiménez-Luna, Sarah Lewis, Victor Garcia Satorras, Bastiaan S Veeling, Regina Barzilay. Tommi Jaakkola, et al. Fast protein backbone generation with se (3) flow matching. arXiv preprint arXiv:2310.05297, 2023a.
Jason Yim, Brian L Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. Se (3) diffusion model with application to protein backbone generation. arXiv preprint arXiv:2302.02277 2023b.
Xinyu Yuan, Zichen Wang, Mareus Collins, and Huzefa Rangwala. Protein structure tokenization: Benchmarking and new recipe. arXiv preprint xxXiv:2503.0008) 2025.
Jiayou Zhang, Barthelemy Meynard-Piganent James Gong, Xingyi Cheng, Yingtao Luo, Hugo Ly, Le Song, and Eric Xing. Balancing locality and reconstruction in protein structure tokenizer. bioRxiv, pp. 2024-12, 2024.
A LLM USAGE STATEMENT
Large Language Models used for polishing grammar, finding typos, and creating plotting code.
13

---

B ETHICS STATEMENT
Al for biology has the potential significantly improve human health, but can also be used for ill. We strongly support the principles outlined in Responsible AI Biodesign
C
METRICS
C.I STANDARD METRICS
This section describes all of the metrics use in assessing model performance throughout the paper. To encourage reproducibility, we release all of these along with any relevant data with our code submission.
We use the Kabsch algorithm, for which there are numerous standard implementations, to align two structures before computing RMSD.
RMSD:
We use biotite's implementation in biotite te. structure. tm_score to compute the TM score between predicted and sampled structure. There are small differences between the biotite implementation and the tmtools implementation, but in our experiments these were on the order of 0.01, so we used the former as it parallelizes more easily.
TM-score:
Designability: We follow standard practice and use ProteinMPNN in Ca only mode to inverse fold eight putative sequences, then use ESMFold fold those into eight output structures. We compute the RMSD/TM score between all eight and report the lowest/highest one.
We compute the TM-score between all pairs of designable structures. challenge with autoregressive models is they produce variable sequence lengths, which can bias diversity results to lower outputs. To mitigate this and provide a fair comparison without needing to generate thousands of structures, we take all pairs within ten residues in size. We report the average over all pairwise comparisons.
Diversity:
We report the average TM score of the closest match in the CATH reference database over all designable structures. We found searching over the full PDB to be computationally too expensive on hardware.
Novelty:
C.2 AUXILIARY METRICS
[ss]RMSD: Alpha helices are highly structured and quite easy for models to reconstruct, beta sheets less so, and coils the most unstructured secondary structure element. This metric com- putes RMSD on proteins that are more than 60% one type of secondary structure provide a more fine-grained view of reconstruction ca- pabilities.
1400
1200
+D
1000
16D
FFPSD
800
600
1260
512D
rFPSD: We compute rFPSD over the CATH dataset by taking features from the trained GearNet fold classifier from Geffner et al. (2025). We have two datasets of features, one from the original CATH dataset and one from the tokenized and reconstructed CATH dataset.
10
102
10'
10a
Number samples
Figure 6: rFPSD converges after 5k samples.
We fit Gaussians to these features using stan- dard estimators and compute the Wasserstein distance as
rFPSD ||,, Me 1/2 TY E -E,,,,,)))
where the subscripts and reference the fit CATH and reconstructions respectively.
FPSD was originally introduced Geffner et al. (2025). A similar metric was introduced in FalFaiings al. 2025). We choose the former over the latter primarily the latter uses the ESM3 structure
14

---

tokenizer, which is trained large number of synthetic structures and carries all of the biases of to- kenization. While not conclusive, the observation that DPLM-2 has a lower rFPSD than ESM3 (but higher RMSD) is encouraging for the notion that we can hill-climb rFPSD to develop tokenizers with better generative capabilities.
Why do only measure rFPSD CATH? The goal of additional metrics provide measurements of quality beyond per-sample quality, which RMSD and TM-score provide, and to correct for biases endowed by model and data choices. The estimators for the Gaussian fitting are
N N
N N - Σ 22 -)(== -
E
For Gaussian noise, both estimators are unbiased and obey the Central Limit Theorem. We empir- ically estimated that required minimum of 5,000 samples to get low variance estimators. We show this in Figure The CATH dataset provides both the diversity and sample count required. We explicitly did want use synthetic data avoid biases like overprediction of alpha helices.
D CRYOET RECONSTRUCTION WITH KANZI TOKENS
This section describes conditional generation task. Cryo-electron tomography (cryoET) is a rapidly advancing imaging technique that creates 3D models of frozen biological samples in a cell (as op- posed to purified protein samples). This allows biologists to understand how proteins and other biological complexes function in their natural state.
cryoET data analysis follows two stages. In the first stage, noisy 2D images are compiled to 3D volume. A number of approaches in the computational imaging literature have approached this problem reconstruction task. In the second stage, user must determine the actual proteins corresponding to a cleaned 3D sample. The input is an L x L x L voxel mask, and the desired output sequence of protein structures L 3. Crucially, no sequence information known, and there may be multiple proteins present (each of unknown length). The standard approach involves searching through extensive proteomic databases and docking proteins using the cross correlation scores to simulated molmap. This is computationally inefficient and quite brittle.
We propose instead treating this as a generative problem, where the objective is autoregressively generate a protein structure conditioned on the voxel cryoET data. A critical insight is that SE(3) invariant tokenizers fail use the conditioning signal (a form of posterior collapse), since they must simultaneously reason about the generated protein and the pose indicated by the input sequence.
Figure 7 schematically depicts both our approach and performance at various protein lengths on a synthetic dataset. For each protein, we simulate a molmap using standard tools, which we then en- code (using either a 3D Unet or surface encoder). This sequence of tokens is provided as in-context conditioning to a language model trained using Kanzi tokens. Even at longer protein lengths, the resulting generated structures have < 3, RMSD agreement with the ground truth protein structure, well within the best docking approaches.
E REPRESENTATION QUALITY
Table evaluates Kanzi tokens a subset of representation quality benchmarks from Yuan et al. 2025 The subset of tasks we consider are residue level tasks designed probe the information capacity of individual tokens. We defer full description of the tasks to Yuan et al. (2025); ) in sum- mary, we include binary classification tasks for binding, catalytic, conserved, epitope, and repeated residues, and regression tasks for measuring flexibility of regions. For each task, take post- quantized representations and use a 2-layer probing MLP. FlexRMSF probing additionally includes sigmoid layer mapping outputs from 0 to
While Kanzi tokens show some representative capabilities, both Kanzi and bio2token underperform invariant models. This is expected, since most residue probing tasks are fundamentally local tasks. The probing layer must leam to convert global residue information to a local estimator, which is challenging for simple two-layer MI.P.
15

---

(a)
(b)
vs Longth (Binned, na per bin)
Uest FSMGD
Encoder
Decoder
2.2
Reselution
2.0
Transformer
-61
-120
<1122
-256
Protein Langth
Figure (a) Schematic our approach to cryoET docking. A volumetric map encoded to sequence of tokens, which are provided as in-context conditioning to language model trained to output Kanzitokens. These tokens are decoded to the original protein. (b) Across simulated resolutions and protein lengths, this approach consistently achieves 3A reconstruction.
bindint
con
ept
rep Family
phys Family
Model
Fold
Family
Fold
Family
Fold
Family
Fold
Family
Fold
Fold
hio2roken
0.489
0.637
0.544
0.539
0.512
0.543
0.514
0.525
0.503
0.569
0.312
0.268
dpim2
0.540
0.794
0.598
0,703
0.570
0.728
0.623
0.700
0.507
0.763
0.475
0.433
csm3
0.497
0.787
0.549
0.749
0.541
0.647
0.614
0.644
0.519
0.677
0.433
0.414
kanzi
0.504
0.692
0.531
0.603
0.525
0.626
0.532
0.604
0.526
0.614
0.298
0.357
Table 4: Updated performance table with revised Kanzi benchmarking. AUROC for all tasks except phys, which uses Spearman's p.
There are several important caveats about these experiments. First, we were unable to directly repro- duce the results in Yuan et al. (2025) using the publicly available checkpoints and codebase. We use the processed datasets provided and train probing layers for 500 epochs with the AdamW optimizer, learning rate 10 44 Despite this, there are still differences in our results for ESM3 compared to Yuan et al. 2025). We will make training/benchmakingg code available upon de-anonymization.
F MODEL
This section contains details on the full model training pipeline.
F.1 HYPERPARAMETERS
We train with AdamW, 31 = 0.9, 32 = 0.95. We train tokenizer models with the following configu- rations (bold indicates base models used for most experiments). We train most models on 1-2 H100s (when training with two GPUs, we use 4 micro-steps for an effective batch size of 256). For each micro step, the elements in the batch are composed single protein with randomly augmented views. This allows us to avoid masking/batching with different sequence lengths.
F.2 DATA PREPROCESSING
We train on the Foldseek-clustered AFDB Drs- We perform the following filtering steps
1. We remove all chains with coil percentage 70%.
2. We remove all chains with mean residue-wise PLDDT < 80.
3. We keep only chains where 80% of the residues have PLDDT 70.
4. We filter all structural homologs of chains in our test sets: a subset of the AFDB (already filtered), CAMEO, CASP14, CASP15, and a subset of CATH.
16

---

Parameter
Value
Batch size
32
Encoder layers
2
Decoder layers
8/12
Attention heads
8
Encoder channels
256/384/512
Decoder channels
256/512/784
Pair-bias channels
64
FSQ levels
(8, 5, 5)
MLP factor
4
Dropout
0.1
Learning rate
1.7 x 10
Warmup iters
1000
LR decay iters
100000
Minimum LR
1 1 10-4
Gradient clip
1.0/None
Sliding window size
4/8/16/None
Micro-steps (grad accum.)
8
QKNorm
True/False
Table 5: Training configuration hyperparameters.
The first two requirements are fairly important to ensure high quality structures. We did not ablate the effect of the last step; various other works like Bose et al. (2023) have imposed similar filtering criteria.
During training, we mean-center all coordinates using Ca position (we do this for full backbone data as well).
easy to see that one cannot have diffusion process that is invariant to translations, since the probability mass would not integrate to (Garcia Satorras et al. (021). It is standard to instead define a diffusion process over a subspace of R L.e. the zero center-of center-of-mase of-mass subspace. The noise interpolation in flow matching is thus well defined, since if e (O, 1) and both have center-of- mass zero, then x, Ix + (1 l)e also does by linearity (Hoogeboom et al. 2022).
We augment proteins with random rotations during training, computed using the following snippet:
from scipy.spatial. transform import Rotation
def sanple_unis mnnifoonrrraatiiissspeetrru
3
return torch.tensort weweetiionrrndonoooooo resshapee
Rotation. randon (andoo()rrd((shspe)).e as_natrix()
F.3 INFERENCE
Diffusion sampling: As mentioned in the main text, we use classifier-free guidance, score anneal- ing, and noise annealing. For the most part, we take the parameters from Proteina without further exploration. A major advantage of the flow autoencoder paradigm is the noise/score scale can be tuned to the particular task hand. Best-of-N sampling: In Table we include single result using best-of-2 sampling. Best-of-N sampling is unique strength of autoregressive models; in each forward pass, we store the bog-likel,hood, and decode the sequence with the best log-likelihood. The fact that we have computationally inexpensive closed-form log- le-likklihood is unique strength of autoregressive models; computing log- ligeeeilihooos requires a full ODE integration for flow models. In other words, the log-like likelihood gives a free estimator for the downstream model quality before running the diffusion process or designability checks. We include this primarily to demonstrate that learned discrete codebook encodes useful knowledge about the token quality.
17

---

(a)
(b)
(c)
109
Seed[O]
Codebook utilizatiinn
Seed[1 137]
An
Hamming distance
Flow loss
Baseline
0.6
10-3
Learned invariance
aLN-sin
No invariance
aLN-shared
0.4
100
0.2
10
0.0
25k
100k
Gradient steps
Gradient steps
Gradient steps
(d)
(e)
(f)
No jitee
Concateeattio
GPT-reg
7
Codebook utilizzation
0.8
Jitter
Cross Attention
Baseline
100
100
GPT prior loss
Flow loss
Flow loss
-2.3A
0.4
0.2
10-1
11A
10-2
50k
25k
50k
25k
50k
Gradient steps
Gradlent steps
Gradient steps
Figure 8: (a) adaLN-single trains unstably, but sharing the weights across all layers matches base- line performance and reduces parameter count by 30%. (b) Two runs at different starting seeds with learned augmentations show very different codebook utilizations. (c) Rotational augmentation in-
creases the Hamming similarity between different views of the same protein, which by default totally different sequences. (d) Codebook jitter mostly just reduces codebook utilization. (e) Cross attention unambiguously makes performance worse. (f) The baseline and GPT-regularized version have identical flow losses, but the GPT-regularized version has an additional per-token cross-entropy loss (the baseline version stays at the random initialization lo0(100))).
A valid concern with best-of-N sampling is mode collapse; the strong diversity metrics in Table 3 emphasize that this is not the case with our results.
G ABLATIONS
This section contains a list of all ablations that we conducted in addition to those presented in the main text. We highlight ablations with positive results in purple, negative results in red, and ambiguous results blue.
Standard diffusion models adaptive layer normalization to condition the model time, but this time conditioning relies on thick linear layers and takes 27% of all model parameters. Subsequent work (Chen et al. 2023) used variant, N-sinN- which uses single adaLN MI.P
AdaLN:
WWaniing For exccution speed, we often relied on the flow loss as proxy for downstream performance. One has to be careful while doing this. The comparison valid and correlates well with reconstruction for identical noise schedules and invariancelequivariann properties change. An encoder that does learn the pose will almost always have higher flow loss than former has better reconstructions.
setups. but this is no longer true when encoder symmetry
that does learn the pose, even reconstructions on about 1k CATH structures. To aid the interested reader. we present flow losses in the plots associated with each ablation and annotate approximate reconstructions. We often omit the actual numerical values of the flow loss, since it is effectively uninterpretahle numerical value (though extremely useful to compare performance
all of the ambiguous cases, we check
experiments).
18

---

on the first layer and has lightweight projections for all subsequent layers. We explore third option, where a single MLP is shared across all layers. We find that while adaLN-single trains very unstably, weight sharing across layers provides equal performance as the standard DiT approach. We add adaLN weight-sharing to Kanzi.
Learned rotational invariance:
To learn rotational invariance, we add an additional random relative rotation between the encoded protein and the encoded noise. This encourages the model to learn an invariant representation, a fact we confirm by observing the Hamming and MSE distances between the learned vectors. These models can train, but are extremely stochastic. Intuitively, the model needs to learn how the conditioning vector can relate in very arbitrary ways to the diffusion process. We ran trials at several random seeds and witnessed very different codebook utilizations. That coupled with the reconstruction performance being worse on CATH by about 0.4A made us reject this change. We exclude learned rotational invariance from Kanzi.
Some work suggests that cross attention may be more effective at conditioning diffusion transformers for more complex conditioning sequences (Chen et al.) (024). We explore conditioning our decoder with cross attention, where the input is a six-eight layer en- coder. Our cross attention layers follow standard practices and attend to the main diffusion trunk following each self-attention layer (Vaswani et al. 2017: Peebles & Xie 2023). This unambigu- ously hurts performance, so exclude it from the model.
Cross-attention conditioning:
Curriculum learning: Since the non-invariant version worked quite well, we thought we could tran- sition slowly between a non-invariant and learned invariant version by gradually augmenting with larger and larger rotations during training (using spherical linear interpolation). This did not help at all and mostly just harmed codebook utilization and flow loss. We exclude curriculum rotational invariance from training.
Explicit rotational invariance: Similar above, it seems reasonable that invariant represen- tation could condition a non-equivariant decoder. As we discuss in Section 4.3 certain invariant encoders like MPNN (Dauparas al., 2022). common choice for protein design tasks, immedi- ately led to codebook collapse and poor performance on par with unconditional flow models. We explored other variants. In particular, we embedded the distances between coordinates di into fea- ture vectors, giving inputs with shape L x L x d. We applied successive layers of row and column attention, then took mean down a sequence axis as the conditioning input. This worked slightly better in that the model didn't 't immediately collapse, but didn't really provide any tangible benefit beyond the aesthetic benefit that different poses would encode to the same protein. We exclude explicit rotational invariance from the model.
Low codebook utilizzaiio: FSQ typically has very high codebook utilization (Mentzer et al. 2023), but our data is low dimensional and highly correlated at input, many values ended up in the same "bin" which led to low usage. We explored a variety of options to increase codebook utilization. As mentioned in the main text, eventually realized that could get high codebook usage by just letting the model train longer, but we leave the discussion here as a potentially useful record of ablations.
Codebook jütter: One strategy to encourage higher codebook use is to "jitter" the codewords. This typically used in VQVAE; for FSQ, we add small noise term before bin quantization. This noise term was quite small but generally made results worse. We exclude codebook jitter from model training.
Rotational transforms: Recently, FFfty et al. (2024) saw improvements in tokenization from applying a generalized rotation before and after the quantization step. This seemed add complexity but have basically no effect on our task specifically, so we exclude rotational transforms from the model.
Non-linear embeddings: We initially started with linear upsampling layer, as is done in Geffne et al.] (2025). Since linear layers preserve correlations in highly structured ways, we went from Linear to Linear -> Swish -> Linear -> LayerNorm. We ablated all of these carefully; the addition of both the non-linearity and the layer normalization both seem important. We add non- linear point-wise encodings to Kanzi.
GPT prior: Prior works have suggested adding a small GPT-regularization term encourage code- books to learn tokens more suitable for autoregressive, weat-token-prediiee modeling (Radford ct al. 2018 2019 Vaswani et al. 2017). We ablate training small GPT model (2 layers, dimension
19

---

Reconstruction (')
CATII
1.010
CASP14
0.835
CASP15
0.909
CAMEO
0.898
AFDB
1.031
Table 6: Long sequence ablation. Sliding window attention with window size 8 maintains strong reconstruction across datasets.
512) tandem with our tokenizers. The cross-entropy GPT loss is added as a regularization term to the flow matching loss. We did not fully ablate these results generative capabilities, but the autoregressive regularization did not harm reconstructions.
Long sequences: To test whether the strong performance of sliding window attention sustains longer sequence lengths, we trained an identical Kanzi model on sequences up to 512 residues. We find that maintain strong performance across all reconstruction benchmarks, see Table
G.1 PAIR-BIASING AND SELF-CONDITIONING
Pair-biasie technique where the attention scores QK" in transformer are biased by projecting learned pair features Zij E Lxxd each layer. These are very common in protein design models since their introduction in AlphaFold2 (Jumper et al.] 2021 Abramson et al. 2024). However, AlphaFold2 and AlphaFold3 both use triangular updates, which many subsequent works forgo due to their computational cost. Thus, it is actually quite unclear if pair-biasing is actually helpful when used without the additional use of triangular attention.
Similarly, self-conditioning (described in Stärk al. 2023) and Chen tions a flow model on previous predictions. Instead of a flow field Vo(x,, 1,x,,,,6 C). we now have flow field xx.x,,xitt.c where X1 is coarse prediction of the clean data from the previ- ous timestep. Self-conditioning is implemented during training by providing coarse predic- tion 50% of the time; the remainder of the time, we simply drop out the prediction (i.e., we do vo (x Ø, é)). Various works have reported improved performance by using self-conditioning.
(2022) condi-
We found it difficult fully disentangle the effects of self-conditioning and pair bias- ing. A complicating challenge was we ulti- mately care about downstream generative per- formance, which reconstruction alone does not track with. For this reason, we conducted a se- ries of ablations updates), self-conditioning, and the two com- bined, all for small 11M parameter model. These results are summarized in Table / and visualized in Figure IO Our primary con- clusion was that while pair-biasing and self- conditioning seemed help reduce both RMSD and TM-score, the effect was not enormous, and our 30M models without either addition were already at 1 resolution or better. Moreover, self- conditioning (possibly due to the large 50% dropout probability) tended to harm codebook utiliza- tion. While codebook utilization is not good per it is useful measure to track. This drop shown in Figure 9 Given these complexifiers, we did not include self-conditioning and op- tionally included pair-biasing in models. However, think exploring self-conditioning and pair- biasing and how they impact generative capabilities, potentially fruitful future avenue of research.
Godabook utilization
0.75
Beselie
0.50
Pair bios
Self cond.
SC
pair-biasing (no triangular
10k
20k
30k
40k
50k
steps
Figure 9: Codebook usage decreases with pair- bias and self-conditioning
20

---

CAMEO
CATH
CASP14
AFDB
1.5
1.5
2.0
1.81 0.7 A
1.52 0.4 A
1.83 0.7 A
1228 0.4
1.5
1.5
Baseline
1.0
0.5
0.5
0.5
0.0
0.0
1.5
1.5
2.0
1.74 x 0.7 A
1.55 0.5
1.67 x 0.5 A
1.14 0.3 A
1.5
Pair blas
1.0
1.0
1.0
0.5
0.5
0.5
1.5
2.0
Self conditoning
1.67 x 0.7 A
1.43 0.5
1.60 : 0.5 A
1.12 = 0.4
1.5
1.0
1.0
0.5
0.5
0.5
0.0
0.0
Pair Pair blas self cond.
1.5
2.0
0.4
1.5
1.0
1.0
Density
1.0
0.5
0.5
0.0
0.0
0.0
0.0
RMSD (A)
RMSD (A)
RMSD (A)
RMSD (A)
Figure 10: Histogram of RMSDs for pair bias and self-conditioning ablations across datasets.
21

---

Method
RMSD
TM-score
CAMEO
1.812
Pair bias
1.741 = 0.671
0,933 + 0.053
Self-cond
1.670 1: 0.679
0.938 0.051
Pair hias Self-cond
1.619 0.602
0,943 0.042
CATH
Bascline
1.519 0.446
0,953 0.027
Pair bias
1.551 +: 0.534
0,947 4 0.037
Self cond
1.431 1: 0.543
0.957 0.030
Pair bias Self-cond
1.420 0.422
0.958 0.028
CASP14
Baseline bias
1.833 0.665
0,948 0.032
1.674 0.465
0.954 0.021
Self-cond
1.596 + 0.473
0.960 + 0.018
Self
AFDB
Bascline
1.282 = 0.421
0.964 + 0.023
Self-comi
1.122 1. 0.356
0.972 1. 0.017
Self-cond
Table 7: Pair-bias and self-conditioning ablation summary statistics (mean std) for RMSD and TM-score across datasets.
H OVERVIEW OF FLOW MATCHING
Flow matching is extensively described elsewhere: see Lipman et al. 2022) or Albergo Vanden-
Eijnden (202) for original presentations and [Esser et al. 2024) for more applied discussion.
Flow matching generalizes diffusion models (Song et al. 2020 Ho et al. 2020), which similarly distributions. While the distinction between flow matching and diffusion is often stated to be straight vs curved paths, this is more of a property of the sampler being used (see Gao et (2024) for a good discussion). As mentioned earlier, for Gaussian flows on IK", we have closed-form expressions to convert between the score and the learned vector field. For this reason, we often use the terms interchangeably. Flow objectives are a good for many molecu- lar diffusion tasks because they avoid the instabilities at early or late timesteps (images solve this problem by clipping, which is not as easy to for points IR$).
interpolate between
For completeness, present the standard flow matching objective and parameterization in notation consistent with the remainder of the paper. Given a data distribution p1 and an easy to sample prior Po. wish to learn family of functions u(x,, Rd x [0, 1] Rd that would push from the distribution pa to Pt. we did have access to a closed-form expression for u(xe, t), then we could simply optimize
LFM = E.- [|| t) ve(x,. +)]
However, do not have such u(x;, t). The critical insight of conditional flow matching is that one can construct conditional path
u(xe Xo. X1, t)p(x; xo, x1, p(x:: 1)
u(xe, Exx-pixxx
In this expression, we have introduced a family of conditional probability paths p(ll:;;0) that in- duces conditional vector fields u[-1,,,,1)) As we do not know p(x, f), we still cannot regress against u(xe, t). However, we define
LCFMs = (Es.e110.1)
Xerxx ()) [[] Xo. x1, -
then by observing that VoLras = VoLCFM. we note that we can optimize LCFM to regress vo ox(x, towards the true vector field. In words, construct a probability path between distri-
22

---

butions, sample a point from each distribution, sample time E [0, 1], and use the aforementioned probability path to sample x, As noted earlier, the standard choice is the linear probability path
pxx, Xp, x = (xxx/txx + (1 t)xo, o²)
where we take = 0 in our experiments (though this is not strictly required, see Gao et al. (0224) This gives rise to the simple vector field u(x: Xo, xI, () = (xI - Xo).
I
COMPARISONS
This section contains further details on the models we benchmark against.
I.1 PRIOR TOKENIZATION LOSSES
First, we describe the objectives prior works have optimized against and discuss the pitfalls of each.
FAPE. The frame-aligned point error (FAPE) loss has become the canonical SE(3)-invariant loss since AlphaFold2 (Jumper al. 2021 Define frame by translation rotation pair T, (Ra, tz). The action a frame on a vector X, is Tix, = Rx, + t.. Given predicted coordinates, predicted frames, true coordinates, and true frames Xj, Tú, x3 Time, the FAPE loss is defined by
FAPE(X,, Ti, the T.Te) IT X, T.rrc X
Every point put into the local reference frame defined by every amino acid. easy to show that this loss is invariant under rotations and translations, but not reflections chirality is important for biomolecules.
FAPE was very impactful, but it has several drawbacks. It scales quadratically between the atom count and the frames. It requires clamping to train stably and can still be very unstable early training. and thus often requires associated binned cross-entropy version. Errors in bond angles can cause kinks in the optimization spectra; AlphaFold2 uses several supplemental bond angle and violation losses help reduce these issues.
dRMSD An alternative to FAPE, used by Hayes et al. (2025), ) is dRMSD, where we regress against inter-atom distances within ground truth and predicted structure. Explicitly, let dij ||x; x,,|2 2|22 Then dRMSD given by
dRMSD ||apped
d 2
dRMSD is invariant under chirality since ||x; - -x,l2 11(-x:) (-xx)]l2. ESM3 thus introduced an additional "backbone vector loss," whose primary purpose seems to be break the chirality symmetry in proteins. Like FAPE, dRMSD is quadratic in atom count, and can struggle to optimize long-range contacts due to being dominated by errors from numerous short range interactions.
RMSD Some works, notably Liu et al. (2024) and Gao et al. 2025), directly optimize the RMSD. This requires performing a rigid alignment using the Kabsch algorithm. This has several issues. First, backpropagating through the Kabsch algorithm is non-trivial and can't 't be done at low pre- cision, which makes efficient training challenging. Second, even at float32 precision, the Kabsch algorithm has discontinuities whenever singular values are particularly close. It also requires deter- minant correction, which can invert the handedness of rotation and introduce another discontinuity. These cause instabilities during training. Third, while these alignments work fine small datasets, they tend to regress to the mean at larger scale, which makes them less suited for modeling more disordered structures and targets of therapeutic interest. Finally, the SVD based alignment is O(L)) in the sequence length L, though this may be faster in practice using iterative approximations to the SVD
Most other losses (such as the multiple losses used in FoldToken) are derivatives or binned versions of one of the three discussed here designed to help mitigate some of these issues.
23

---

1.2 MODELS
ESM-3 We pull the model as described at https://github. comm/eeolutiooaryscaleeeesss Sev- eral others have observed that ESM3, having been trained on a large amount of metagenomic data, struggles to generate designable sequences, Geffine et al. (2025) reports designability of 22%. We find that
improve this slightly by increasing the number of steps from the default 8 on the
GitHub up to single token per step. We hypothesize this is partially because the ESM3 tokenizer is so high variance; even at high designabilities, the mean scRMSD remains high due to few decoded structures with incredibly high RMSDs (>20 Å).
ESM-AR ESM-AR was trained from scratch on a mix of AFDB data and PDB data. has dimension 1280, MLP dilation factor 2, 16 layers, and 16 heads. Otherwise, it is standard transformer with pre-LayerNorms and GELU activation functions. ESM-AR seems to improve ESM3's designability performance, but does not continually scale with techniques like best-of-N sampling, which we again attribute to the aforementioned high variance decoding in the ESM3 decoder.
and oth- erwise follow the generation example exactly using generate_dplm2 py. A potential issue we noted was while our generations had similar scRMSDs as those reported in Wang al. (2024), we slightly underperformed the reported scTM values. The difference of z 0.14 seems large, but given that scRMSD is generally acknowledged to be a more accurate reflection of protein structure that TM, seems within reason.
DPLM2 We pull the DPLM2 weights/code from https: / thub. www. wwt.//oeeaann////pppppoooaa
The D-LM-2 structure tokenizer is fine-tuned on 220k structures sourced from the PDB SwissProt. However, the structure encoder a large pretrained GVP-Transforr (Hsu et al. 2022) trained on 12 million AFDB structures. We thus include these in the DPLM-2 data count for Figure
DPLM-AR The DPLM-2 codebook has size 8192, compared to the 4096 token ESM codebook.
Other than the projection and embedding layers, the DPLM-AR model is identical the ESM-AR model. Unlike the ESM-AR model, the DPLM-AR model slightly underperforms DPLM2.
IST We used the 1.7k codebook tokenizer, but were unable to generate any des- ignable structures using the generative models released with the InstaDeep Structure To- kenizer. Otherwise, we exactly followed the encoding/decoding process described in https://github https//gggttbb..omm/nnntteeeepp
bio2token We pulled weights and scripts from http:: www htmm///laghhpppioeerrn/noooooooo As bio2token operates over point clouds, rather than frames, we tokenized and reconstructed Co and full backbone coordinates separately. We excluded bio2token from CATH comparisons as explicitly trained on the entirety of CATH.
Fole'Token4 We pull FoldTooken from https: / www. comm/44BooFolloooen _open and follow the described reconstruction process. The language model FoldGPT described in the original paper does not have publicly available weights.
Other tokenizers There are three additional tokenizers we do not explicitly benchmark against, de- scribed in Zhang et al. (2024). Zhang et al. (2024), and [in et al. (2023). The former has significant GPU and RAM requirements that made it infeasible to do thorough benchmarks on A100s. None of the three have public models with generative capabilities, so we opt to focus on tokenizers that demonstrate downstream performance along generative axes and provide good coverage of different tokenization techniques.
J EXTENDED RESULTS
This section contains the tables from the main text with additional error bars. Our primary obser- vation is that transformer based tokenizers to have number of catastrophic that lead to very high variances in the output. We will make all these experiments public along with our benchmarking repo upon acceptance and deanonymization.
24

---

CAMBO
CASP14
CASPI5
CATH
AFDB
RMSD (U)
(1)
RMSD
TM
RMSD
TM
RMSD
TM
RMSD
TM
DPP M2
1.30
0.14
1.33
0.951
1.022
2
0.866 1
1641 1.27
0.807
000771000 0.00
4.576.66:04
1..15
ESMB
0.00
0.49
0.462 0.33
0,987
4.002
1,21
1,9
1:0.5
0.977
+: 040
2.384
0.915:0.111
2.5
0.881 0.17
2.04
2.90
0:08
666948861
0744
1.
1.28
1.57
0.900
0.6
2.161114
0.85
IST
1.67
0.916=0.11 0.13
0,900
0.960
0
0.953
+ 0.00
1.201
0.72
0,940
0.11
3.672
0.862 0.11
buo2bokea
0.24
0:02
0.999
0.993
0:542
004
0.902
Keroi (30M)
0.93610027
0.94
1:11
0.09
0.
0.5%
0:540
004
0.947 1.449
Kanzi (30M)*
0.817 0.29
0.960
0.03
0.19
0.07
1,267 0.35
0963::0011
0.953 0.57
0,955
0.04
0.962 0.339
0.02
Kanzi (1IM)
10016+00.3
0.337
+ 0.05
0.11
0,954 + 0.01
+.33
0.955 + 0.01
1.156 0.89
0,924
+ 0.04
1.210
0.69
0,934 + 0.004
(1IM)
0.863 0.225
0:10
0908
0.
0.994
0:950
0:466
0.952
Table 8: Reconstruction metrics across tokenizers for Ca reconstruction with errors. Smaller datasets unsurprisingly have very large emror bars. An advantage of flow tokenizers is they tend to have lower variances across the board. Other tokenizers seem to have a small number of coordi- nates with catastrophically high RMSDs, which leads really large standard deviations.
CAMPO
CASP14
CASPIS
CATH
AFDR
RMSD (4)
(1)
RMSD
RMSD
RMSD
DPLM2
1831 1.39
0.39
0.959 +
2.146
1.75
0.93
1.777+1.88
0,925
0.8
4.646++6.05
0.880
ESM3
0 1 1.60
0.900 - 0.05
0.463 4 0.33
0.994 0.00
1018 1:
1.87
0
0.06
1 2.200
0.971
1006
2,378
4.088
0.944
1.11
FoldToken
2 3:006
0902200.09
2.523 3.01
0.95
6.58
0.809
0.22
1
0.94
0.05
2,120 1.44
0.907
4 0,08
1.6361 1.84
0.054
D.896
0.21
0:01
1.24
0.975.0000
0.7
0.956
1 0.04
1.859
0.904
1.111
boo2token
1.089 =: 0.28
0.963
= 0.04
0.998 à 0.24
0.966
1.367
0.960 0: 0.04
0 =: 0.25
0.958
:: 0.03
1,201
0:09
0.949 0.003
Kinsi (30M)*
0.33
0,973 0.03
0.20
0.981 0.001
0
0.000
0.972 0.074
0.02
1.165 1,0.0
0.969 0.02
Table 9: Reconstruction metries across tokenizers for full backbone reconstruction with errors.
K
EXTENDED THOUGHTS
This section contains extended discussion of several non-critical points raised in the main text. We welcome further discussion on any of the following.
K.1 CONNECTIONS BETWEEN FLOW MODELS AND THE STRUCTURE DECODER
On first pass, it is somewhat surprising that flow autoencoders are so parameter efficient. The origi- nal structure module Jumper et al.) (202T) actually reused weights for each iteration of the structure decoder. This required careful tuning and several techniques to prevent instability during training. Subsequent methods that use the structure decoder a generative or tokenization context generally opt not to share weights (Gaujac et al., 2024] Yim et al. 2023b). The structure decoder can be thought of as a very coarse diffusion model with only steps. where the prior is the identity rota- tion. Our transformer generalizes this to take an arbitrary number of steps. From this perspective, the performance of diffusion transformers (both in our work and in models like AlphaFold3) makes much more sense.
K.2
WHY CAN POINTWISE LAYERS RECONSTRUCT SO WELL?
We found intriguing that even very small window sizes with sliding window attention (SWA) gave good reconstructions, so we ablated a simple pointwise MLP. These ablations are discussed in Section 4.3 While these models underperformed our models with SWA, they nonetheless gave re- constructions of around 1.3 A. However, generative models trained over these codebooks performed very poorly. Why is this?
The purpose of the encoder in an image tokenizer like VQGAN is to pool local information to create meaningful representations. This has both a semantic and a technical motivation. Semantically, the encoder learns local components of an image, and powerful attention mechanisms can then leam long-range interactions. Technically, the downsampling reduces the sequence length of the transformer manageable size
Because our protein structure tokenizers operate on raw coordinates, there is already some long- range information present, since (roughly speaking) proteins near the start or the end of the chain
25

---

will often have bigger numerical values than near the center. Since this information is already present, the model just needs to learn to transform in a way that's more amenable for sequence modeling, which is not that hard. Because generative models explicitly operate sequences of tokens, however, this representation is not necessarily amenable to generation.
K.3 THE CASE FOR AUTOREGRESSIVE GENERATION
Most models the protein generation space are diffusion or discrete diffusion models, with good reason. These models work really well. They have saturated designability over the past year and they're proven to create proteins that can be designed in a lab. One could be forgiven for thinking that the focus on autoregressive models is unnecessary.
However, a major benefit of autoregressive models that sequence lengths can be learned. The use of diffusion models seems to be heavily entrenched into the protein folding/inverse folding problem, where the sequence length is known a priori. When a model folds a sequence into a structure, by design it knows the structure size. Being able to generate protein sequences of variable lengths is a critical skill as we move towards more diverse tasks. The canonical example in situ modeling, where cannot purify single protein the sequence length is unknown. In principle, one could sample a bunch of different lengths with a diffusion model and score them, but this comes with its own of issues. A human operator would need make this decision beforehand and results would likely be quite brittle (e.g., the difference between having two and three proteins at inference time would be substantial). An autoregressive model could just learn to generate proteins of appropriate size a priori.
There are other examples (e.g., binding to membrane proteins, which is particularly important for therapeutics). With enough priori information, diffusion model generally outperforms an au- toregressive model because it has full bidirectional communication. This gives diffusion models significant advantage in benchmarks like designability, where people typically just generate proteins of different sizes. However, we generally think that as we move to more complex, in-the-wild tasks, giving models the capability to reason about protein size will be valuable capability.
26


Reviews:
[{"review_id": "T0mJK5HcHb", "reviewer": "Reviewer_6hfk", "summary": "The authors present a flow-based protein structure auto-encoder (\"Kanzi\") useful as a protein tokenizer. The primary contribution is that, as far as I am aware, it is the only example of flow-based structure auto-encoder. Secondary to that, but still very important, is that the encoder and decoder are not invariant or equivariant, continuing the recent trend to show that such complex models are not required to adequately model proteins. The authors additionally use the model to generate novel structures.", "questions": "The authors might want to cite Ellmen, ... Deane \"Transformers trained on proteins can learn to attend to Euclidean distance\", which includes an interesting discussion of how attention mechanisms process real space coordinates.", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The authors presentation is mostly very clear, and it does a good job of discussing the previous literature. They present a novel approach to the problem of structure tokenization. I particularly encouraged to see that the output is simply the backbone atomic coordinates, which should greatly simplify the tokenizer's use. It is intriguing that the model uses real space coordinates _and_ a sequence id embedding; see questions for more on this.\n\nI like that the authors show that the resulting tokenizer can be used to generate designable structures. But see weaknesses for additional comments on this.\n\nClearly Kanzi is successful: the 30M parameter version trained with \"under optimized\" hyperparameters is best or second best across most metrics, covering the core set of structure prediction data, while being _much_ smaller. \n\nThe discussion of ablations is especially useful, although I do wish it were more detailed.", "weaknesses": "This is a small weakness but the authors seem to flip between using \"diffusion\" and \"flow matching\" to describe their approach. It seems that flow matching better matches what they're doing--either that or there is something that needs to be clarified substantially in their writing. In any case diffusion and flow matching are not exactly the same thing, so the authors should be precise.\n\nThere are two significant weaknesses in the paper that I'd like to see addressed:\n1. The authors should examine how their tokenizer performs at other downstream tasks that are relevant for protein language models. Many are detailed in a recent paper from ICML 2025: Xinyu Yuan et al., Protein structure tokenization: Benchmarking and new recipe.\n\n2. While the authors show that their auto-regressive structure generator is capable of generating designable structures using the conventional definition of generating a sequence which folds back to the same structure, this is highly dependent on the particular folding model used. The authors chose ESMFold, which I found curious given they train on a sample of AFDB. The authors should comment on this choice given that better models are available.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2aNaUqGzOo", "reviewer": "Reviewer_xSwq", "summary": "This paper presents Kanzi, a flow-based protein structure tokenizer that replaces SE(3)-invariant components with a simpler diffusion autoencoder trained using a flow matching loss. The approach removes the need for complex invariant losses and geometric attention while maintaining or even improving reconstruction performance compared to prior tokenizers like ESM3, DPLM2, and FoldToken. \n\nKanzi is then used to train an autoregressive model (Kanzi-AR) for structure generation, demonstrating competitive designability and diversity.", "questions": "1. Compare with two more baselines: \n- Cheap [1]\n- AminoAseed [2]\n\n2. Add representation quality evaluation. For example, the benchmark from AminiAseed [2].\n\n3. Questionable use of diffusion sampling in AR models. Table 3 claims Kanzi-AR uses inference-time sampling tricks (Eqn. 5), which conceptually apply to diffusion-based samplers, not discrete autoregressive decoders. As I can understand, AR simply uses the discrete structure tokens for autoregressive training and sampling. There should not be any diffusion sampling process involved? Clarification is needed.\n\n[1] Tokenized and continuous embedding compressions of protein sequence and structure \n\n[2] Protein Structure Tokenization: Benchmarking and New Recipe", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- **The experimental setup is solid and shows careful design choices** — e.g., diverse test datasets, detailed RMSD/TM metrics, and fair consideration of computational trade-offs. \n\n- **The writing is clear, and figures are well-designed.** The methodology section carefully explains the training and inference setup. \n\n- **Simplified yet effective formulation.** The use of a single flow matching loss instead of multiple SE(3)-invariant losses reduces training complexity and improves stability. This work is more like an altogether engineering refinement than a conceptual breakthrough. The proposed simplification aligns with recent diffusion autoencoder trends in vision. Though the novelty in biological context is limited, the \"making it simple and scalable\" is indeed a very important next step for this field.", "weaknesses": "- **Performance is not that out-standing**: The reconstruction performance and generation performance do not look that out-standing. ESM3 seems to still be the best of all. Inference-time sampling tricks from Proteina seems to be able to boost Kanzi-AR to a next level, though unfortunately neither ESM3 or DPLM2 used this trick. This might be an unfair comparison.\n\n- **Only evaluating on reconstruction and generation, missing representation quality**: From image domain, there are some papers discussed about one point: not necessarily the best reconstruction quality leads to a better tokenizer. The representation quality also matters. e.g., see the table 1 in RAE paper (https://arxiv.org/pdf/2510.11690v1) for MAE-B and DINOv2-B. And there is a benchmark designed for structure tokenizer representation quality: from AminoAseed [1] paper. Highly suggest to also benchmark representation quality rather than reconstruction quality alone.\n\n- **Missing comparisons to other structure tokenizer benchmarks**: see questions.\n\n- **Minor typos**: There are also minor typographical issues (e.g., lines 794–797 contain “¿” and “¡” characters).\n\n\n[1] Protein Structure Tokenization: Benchmarking and New Recipe", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "i3soMTtMna", "reviewer": "Reviewer_zKJZ", "summary": "This paper introduces Kanzi, a non–SE(3)–equivariant diffusion autoencoder–based protein structure tokenizer, trained using a single flow-matching objective on global 3D coordinates. The authors argue this architecture simplifies training relative to SE(3)-invariant tokenizers and demonstrate strong reconstruction and unconditional structure generation performance compared to recent tokenizers such as ESM3, DPLM2, FoldToken, and IST.", "questions": "1. Does rotating a protein structure change the tokenization output? Have you evaluated rotational consistency quantitatively?\n\n2. Why is retrieval / homology search omitted? Do Kanzi tokens perform poorly under similarity search?\n\n3. Can Kanzi support conditional generation (e.g., topology, motif, scaffold constraints)? If yes, why is it not reported?\n\n4. You do not perform residue-level local centering or relative-frame coordinate normalization (e.g., per-residue local frame / backbone-centric coordinates) — without such normalization, how are the learned tokens supposed to be interpretable or physically meaningful?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Clear motivation to simplify protein structure tokenization by removing architectural and loss engineering complexity.\n\n- Technically solid implementation using flow-matching and FSQ-based codebooks.\n\n- Strong reconstruction quality across multiple protein benchmarks despite relatively small model size and training data.\n\n- Interesting empirical findings: non-equivariant encoders outperform equivariant ones under flow objectives, codebook utilization emerges late in training.\n\n- Introduction of rFPSD as a distributional reconstruction metric is potentially valuable.", "weaknesses": "1. No evaluation of SE(3) stability / invariance — critically, the tokenizer may output different tokens for rotated versions of the same protein, which invalidates many downstream use cases (retrieval, clustering, homology, interpretability). This is not even measured.\n\n2. No retrieval / similarity / StructTokenBench[1]-style evaluation, even though retrieval is a core purpose of structured tokenization (cf. FoldSeek[2], ).\n\n3. No conditional generation experiments, despite stating generative capability as a core contribution; all reported results are unconditional only.\n\n4. Claims of “smaller model rivaling ESM3/DPLM2” are not apples-to-apples — those are protein sequence / multimodal models, not structure-only.\n\n5. Scope of “tokenizer quality” is too narrow — heavily focused on reconstruction, insufficient multi-dimensional evaluation (e.g. sensitivity, explainability, stability, controllability).\n\n[1] Protein Structure Tokenization: Benchmarking and New Recipe\n\n[2] Fast and accurate protein structure search with Foldseek", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WLU4tYg6Na", "reviewer": "Reviewer_BvV7", "summary": "They build a discrete VQVAE based tokenizer for protein structures, where the encoder maps 3D coordinates to discrete codes and the decoder is a DiT-based flow matching model. They additionally train an autoregressive prior over the latent codes for generation. The approach, Kanzi, simplifies aspects of prior approaches, including the use of a diffusion loss (as opposed to prior SVD-based losses), replacing frame representations with global coordinates, and using standard attention instead of SE(3) invariant architectures. Kanzi outperforms prior token-based structure generative models and is more parameter efficient.", "questions": "- Do you think the sliding window attention would still suffice for longer sequences?\n- While the model achieves superior reconstruction and competitive generation for token‐based methods, it still lags continuous diffusion models. What are the main causes of this gap?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "- The authors perform extensive benchmarking against prior approaches on both reconstruction and generative tasks. The results demonstrate that the model is competitive against the prior SOTAs on both tasks.\n- They perform a series of ablations on encoder variants (invariant vs not invariant), attention window sizes, and model size.\n- Simplifying the auto-encoder loss is pretty significant, as it significantly reduces the cost of training from O(L^3) or O(L^2) to O(L).\n- Using a diffusion decoder also allows you to use classifier-free guidance and the diffusion noise scales to balance the tradeoff between diversity and sample quality.", "weaknesses": "- The idea isn't incredibly novel. The paper is a combination of many design choices (diffusion decoder, DiT rather than SE(3) invariant attention, FSQ discretization) rather than a single great idea. \n- You mentioned the cost of loss functions used in prior works. Maybe you can do an experiment measuring the iteration speeds and memory scaling of each approach? Note also that diffusion model autoencoders have a multiplicity (number of diffusion timesteps per step), which can increase the memory use / runtime of this approach. Also, despite these supposed gains, you only trained on proteins of length <256. An experiment scaling to larger systems may be worthwhile.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)