Masked Autoencoders Are Effective Tokenizers for Diffusion Models
Authors: Hao Chen
Yujin Han
Fangyi Chen
Xiang Li
Yidong Wang
Jindong Wang
Ze Wang
Zicheng Liu
Difan Zou
Bhiksha Raj
Venue: ICML 2025 spotlightposter (2025)
Keywords: Diffusion Models
Autoencoders
Masked Autoencoders
Primary area: deep_learning->generative_models_and_autoencoders

Abstract: Recent advances in latent diffusion models have demonstrated their effectiveness for high-resolution image synthesis. However, the properties of the latent space from tokenizer for better learning and generation of diffusion models remain under-explored. Theoretically and empirically, we find that improved generation quality is closely tied to the latent distributions with better structure, such as the ones with fewer Gaussian Mixture modes and more discriminative features. Motivated by these insights, we propose MAETok, an autoencoder (AE) leveraging mask modeling to learn semantically rich latent space while maintaining reconstruction fidelity. 
Extensive experiments validate our analysis, demonstrating that the variational form of autoencoders is not necessary, and a discriminative latent space from AE alone enables state-of-the-art performance on ImageNet generation using only 128 tokens. MAETok achieves significant practical improvements, enabling a gFID of 1.69 with 76× faster training and 31× higher inference throughput for 512×512 generation. Our findings show that the structure of the latent space, rather than variational constraints, is crucial for effective diffusion models. Code and trained models will be released.

Masked Autoencoders Are Effective Tokenizers for Diffusion Models
Hao Chen 12 Yujin Han" Fangyi Chen
Xiang Li!
Yidong Wang" +
Jindong Wang S Ze Wang?
Zicheng Liu?
Difan Zou
Bhiksha Raj*
Abstract
Although originally formulated in pixel space (Ho al.. 2020; Dhariwal & Nichol, 2021), subsequent research has shown that operating in a latent space compressed repre- sentation typically learned by tokenizer - can substantially improve the efficiency and scalability of diffusion models (Rombach et al., 2022a). By avoiding the high- ciabdimmeniinal pixel domain during iterative diffusion and denoising steps, latent diffusion models dramatically reduce computational overhead and have quickly become the de facto paradigm for high-resolution generation (Esser et al., 2024).
Recent advances in latent diffusion models have demonstrated their effectiveness for high- resolution image synthesis. However, the prop- erties of the latent space from tokenizer for bet- ter learning and generation of diffusion models remain under-explored. Theoretically and empiri- cally, we find that improved generation quality closely tied to the latent distributions with better structure, such the ones with fewer Gaussian Mixture modes and more discriminative features. Motivated by these insights, we propose MAE- Tok, an autoencoder (AE) leveraging mask mod- cling to learn semantically rich latent space while maintaining reconstruction fidelity. Extensive ex- periments validate our analysis, demonstrating that the variational form of autoencoders is not necessary, and a discriminative latent space from AE alone enables state-of-the-an performance on ImageNet generation using only 128 tokens. MAETok achieves significant practical improve- ments, enabling gFID of 1.69 with 76x faster training and 31x higher inference throughput for 512x512 generation. Our findings show that the structure of the latent space, rather than varia- tional constraints, is crucial for effective diffusion models. Code and trained models are released!.
However, a key question remains: What constitutes a "good" latent space for diffusion? Early work primarily employed Variational Autoencoders (VAE) (Kingma, 2013) as tok- enizers, which ensure that the learned latent codes follow relatively smooth distribution (Higgins et al., 2017) via Kullback-Leibler (KL) constraint. While VAEs can em- power strong generative results (Ma et al., 2024; Li et al., 2024b; Deng et al., 2024), they often struggle to achieve high pixel-level fidelity in reconstructions due to the im- posed regularization (Tschannen et al., 2025). contrast, recent explorations with plain Autoencoders (AE) (Hinton Salakhutdinov, 2006; Vincent al., 2008) produce higher- fidelity reconstructions but may yield latent spaces that insufficiently organized or too entangled for downstream generative tasks (Chen et al., 2024b). Indeed, more re- cent studies emphasize that high fidelity to pixels does not necessarily translate into robust or semantically disentan- gled latent representations (Esser et al., 2021; Yao & Wang, 2025); leveraging latent alignment with pre-trained models can often improve generation performance further (Li et al., 2024c; Chen et al., 2024a; Qu et al., 2024; Zha et al., 2024).
1. Introduction
Diffusion models (Sohl-Dickstein et al., 2015a; Ho et al., 2020; Rombach et al., 2022a; Peebles & Xie, 2023) have recently emerged as a powerful class of generative models, achieving state-of ffate-of-the-ar (SOTA) performance on various image synthesis tasks (Deng et al., 2009; Ghosh et al., 2024).
In this work, we attempt to answer this question by investi- gating the interaction between the latent distribution learned by tokenizers, and the training and sampling behavior of diffusion models operating in that latent space. Specifically, we study AE, VAE and the recently emerging representation aligned VAE (Li et al., 2024c; Chen et al., 2024a; Zha et al., 2024; Yao & Wang, 2025), by fitting Gaussian mixture model (GMM) into their latent space. Empirically, we show that a latent space with more discriminative features, whose GMM modes are fewer, tends to produce a lower diffusion loss. Theoretically, we prove that latent distribution with fewer GMM modes indeed leads to a lower loss of diffusion
Equal contribution 'Carnegie Mellon University AMD 'The University of Hong Kong *Peking University 'William & Mary. Correspondence to: Hao Chen <haoc3 @andrew.cmu.ed
Proceedings of the nnnd International Conference on Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
pttppt///gtttee www..thh//ooooo..t http///hhhhhaaa/hceeoaa cont wnwousntoootooooooo tokenttttkkeeeee

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
Figure 1. Diffusion models with MAETok achieves statc-of-the-art image generation on ImageNet of 512x512 and 256x 256 resolution.
models and thus better sampling during inference.
ing the AE's encoder, thus preserving its well-organized la- tent space, and fine-tuning only the decoder, we can recover strong pixel-level reconstruction fidelity without sacrificing the semantic benefits of the learned representations.
Motivated by these insights, demonstrate that diffusion models trained on AEs with discriminative latent space are enough to achieve SOTA performance. We propose to train AES as Masked Autoencoders (MAE) (He et al., 2022; Xie et al., 2022; Wei et al., 2022), a self-supervised paradigm that can discover more generalized and discriminative rep- resentations by reconstructing proxy features (Zhang al., 2022). More specifically, we adopt the transformer architec- ture of tokenizers (Yu et al., 2021; 2024c; Li et al., 2024c; Chen et al., 2024a) and randomly mask the image tokens at the encoder, whose features need to be reconstructed at the decoder (Assran et al., 2023). To maintain a pixel de- coder with high reconstruction fidelity, we adopt auxiliary shallow decoders that predict the features of unseen tokens from ones to learn the representations, along with the pixel decoder which is normally trained previous tok- enizers. The auxiliary shallow decoders introduce trivial computation overhead during training. This design allows us to extend the MAE objective that reconstructs masked image patches, to simultaneously predict multiple targets, such as HOG (Dalal Triggs, 2005) features (Wei al., 2022), DINOv2 features (Oquab et al., 2023), CLIP embed- dings (Radford et al., 2021; Zhai et al., 2023), and Byte-Pair Encoding (BPE) indices with text (Huang et al., 2024).
Extensive experiments on ImageNet (Deng et al., 2009) demonstrate the effectiveness of MAETok. It addresses the trade-off between reconstruction fidelity and discriminative latent space by training the plain AEs with mask modeling, showing that the structure of latent space is more crucial for diffusion learning, instead of the variational forms of VAEs. MAETok achieves improved reconstruction FID (rFID) and generation FID (gFID) using only 128 tokens for the 256x 256 and 512x512 ImageNet benchmarks.
Our contributions can be summarized as follows:
We establish a connection between latent space structure and diffusion model performance through both empirical and theoreti- cal analysis. We reveal that structured latent spaces with fewer Gaussian Mixture Model modes enable more effec- tive training and generation of diffusion models.
Theoretical and Empirical Analysis:
We train plain AEs using mask modeling and show that simple AEs with more discriminative latent space empower faster learning, better generation, and higher throughput of diffusion models, showing that the variational regularization of VAE not necessary.
MAETok:
Furthermore, we reveal an interesting decoupling effect: the capacity to leamn discriminative and semantically rich la- tent space at the encoder can be separated from the capacity to achieve high reconstruction fidelity at the decoder. par- ticular, higher mask ratio (40-60%) in MAE training often degrades immediate pixel-level quality. However, by freez-
SOTA Generation Performance: Diffusion models of 675M parameters trained on MAETok with 128 tokens achieve performance comparable to previous best models on 256 ImageNet generation and outperform previous 2B USiT at 512 resolution with 1.69 gFID and 304.2 IS.
2

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
Fewer modes are sufficient adequately represent the la- tent space distributions of VAVAE compared to those of AE and VAE, highlighting simpler global structures in its latent space. Correspondingly, Fig. 2b reports the training losses of diffusion models with AE, VAE, and VAVAE, which (al- most) align with the GMM losses shown in Fig. 2a, where fewer modes correspond to lower diffusion losses and better gFID. This alignment validates our intuition, confirming that latent spaces with fewer modes and thus more separated and discriminative features can reduce the learning difficulty and lead to better generation quality of diffusion models.
100
VAE VAVKE 30 MANTTA
VaMAE (2R1D-1258))
(25M NILL
Dilfnsion Looe
30
2.0
3.4
00
150
2000
230
300
0.2
1001
100
Comporente
Training Stops
(a) GMM Loss
(b) Diffusion Loss
Figure 2. GMM fitting on latent space of AE, VAE, VAVAF, and MAETok. Fewer GMM modes in latent space usually corresponds to lower diffusion losses and better generation performance.
Theoretical Analysis. After observing experimental phe- nomena that align with our high-level intuition, we further present a concise theoretical analysis here to justify the ra- tionale behind it, with more details provided in Appendix A.
2. On the Latent Space and Diffusion Models
To study the relationship of latent space for diffusion models, we start with popular tokenizers, including AE (Hinton & Salakhutdinov, 2006). VAE (Kingma. 2013), representation aligned VAE, i.e., VAVAE (Yao & Wang, 2025). We train our AE and VAE tokenizers under the same training recipe and the same dimension for fair comparison. We train diffusion models on them and establish connections between latent space properties and the quality of the final image generation through empirical and theoretical analysis.
Following the empirical analysis setup, we first consider latent data distribution in d dimensions modeled as a GMM with K equally weighted Gaussians:
1 K
Po = K W I),
(1)
Considering the classic diffusion model DDPM (Ho et al., 2020) and following the training objective as Shah et al. (2023), the score matching loss of DDPM at timestep / is
Empirical Analysis. Inspired by existing theoretical work (Chen et al., 2022: 2023; Benton et al., 2024), our investiga- tion of the connection between latent space and generation quality starts with high-level intuition. With optimal dif- fusion model parameters, such as sufficient total time steps and adequately small discretization steps, and with assumed similar capacity of tokenizer decoders, the generation qual- ity of diffusion models, i.e., the learned latent distribution, is dominated by the denoising network's training loss (Chen et al., 2022; 2023; Benton et al., 2024), while the effec- tiveness of training diffusion model via DDPM (Ho et al., 2020) heavily depends on the hardness of learning the latent space distribution (Shah et al., 2023; Diakonikolas et al., 2023; Gatmiry et al., 2024). Specially, when the training data distribution is too complex and molti-modal, i.e., not discriminative enough, the denoising network may struggle to capture such entangled global structure of latent space, resulting in degraded generation quality.
min [E]| sw (x, t) - Vx logg)((xx)] w
(2)
where sw () represents the denoising network and Vx log (x)) denotes the oracle score function.
Then, we establish the following theorem to show that more modes typically require larger training sample sizes for diffusion models to achieve comparable generation quality.
Theorem 2.1. (Informal, see Theorem A.7) Let the data distribution be mixture of K Gaussians defined in Eq. (1). Then assume the norm of each mode is bounded by some constants, let d be the data dimension, T be the total time steps, and be a proper target error parameter. In order to achieve a O(T?) error in KL divergence between data distribution and generation distirbution, the DDPM algorithm may require using n n' number of samples:
Building upon this intuition, we use the Gaussian Mixture Models (GMM) to evaluate the number of modes in alterna- tive latent space representations, where a higher number of modes indicates a more complex structure. The details of GMM training are included in Appendix B.3, Fig. 2a ana- lyzes the GMM fitting by varying the number of Gaussian components and comparing their negative log- likeliill losses (NLL) across different latent spaces, where a lower NLL indicates better fitting quality. We observe that, to achieve comparable fitting quality, i.e., similar GMM losses, VAVAE requires fewer modes compared to VAE and AE.
K Ktd=Ba
n' = 0
E2
where the upper bound of the mean norm satisfies max, ||p||| < B.
Theorem 2.1 combines Theorem 16 from (Shah et al., 2023) and Theorem 2.2 from (Chen et al., 2023), showing that to achieve a comparable generation quality O(Te)), latent spaces with more modes (K) require a larger training sam- ple size, scaling as O (K*)). J.This theoretically help explain
3

---

Masked Autocncoders Are Effective Tokenizers for Difflusin Models
Autoencoder
Takan
Decoder. To reconstruct the image, use a set of N learnable image tokens R We concatenate these mask tokens with h as the input to the decoder, and takes only the outputs from mask tokens for reconstruction:
Floo Reyro
Enoo Mock Tokssw
Uvd. Irage loward
VIT Encoder
VIT Decoder P
x = D((e;h]))
(5)
We then use a linear layer on top of & E RN to regress the pixel values and obtain the reconstructed image I.
Auxliary Decoders
Position Encoding. To encode spatial information, we ap- ply 2D Rotary Position Embedding (RoPE) to the image patch tokens x at the encoder and the image tokens at the decoder. In contrast, the latent tokens z (and their encoded counterparts h) use standard 1D absolute position embed- dings, since they do not map to specific spatial locations. This design ensures that patch-based tokens retain the notion of 2D layout, while the learned latent tokens are treated as a set of abstract features within the transformer architecture.
Other Targats
Aux. Decoder Diur
Aux. Decoder Daur
HOG Features
DINOVZ Festures
Figure 3. Model architecture of MAETok. We adopt the plain ID autoencoder (AE) tokenizer, with vision transformer (ViT) en- coder E and decoder D. MAETok is trained using mask modeling encoder, with mask ratio of 40-60%, and predict multiple get features, c.g., HOG, DINO-v2, -2, and CLIP features, of masked tokens from the unmasked ones using auxiliary shallow decoders,
Training objectives. We train MAETok using the standard tokenizer losses as in previous work (Esser et al., 2021):
why, under finite number of training samples, latent spaces with more modes (e.g., AE and VAE) produce worse genera- tions with higher gFID. We provide additional experimental results in Appendix A, demonstrating that these latent distri- butions share comparable upper bounds B, thus justifying our focus primarily on the impact of mode number K.
L = Crecom + X Ppeccep X2 Ladv;
(6)
with Lrecon Lepecceps and Lad, denoting as pixel-wise mean- square-error (MSE) loss, perceptual loss (Larsen et al., 2016; Johnson et al., 2016; Dosovitskiy & Brox, 2016; Zhang et al., 2018), and adversarial loss (Goodfellow et al., 2020; Isola et al., 2018), respectively, and x and X2 being hyper- parameters. Note that MAHTok is plain AE architecture, therefore, it does not require any variational loss between the posterior and prior as in VAEs, which simplifies training.
3. Method
Motivated by our analysis, we show that the variational form of VAES may not be necessary for diffusion models, and simple AEs are enough to achieve SOTA generation perfor- mance with 128 tokens, as long as they have discriminative latent spaces, i.e., with fewer GMM modes. We term our method as MAETok, with more details as follows.
3.2. Mask Modeling
Token Masking at Encoder. A key property of MAETok is that we introduce mask modeling during training, following the principles of MAE (He et al., 2022; Xie et al., 2022), to learn a more discriminative latent space in a self-supervised way. Specifically, we randomly select a certain ratio, e.g., 40%-60% of the image patch tokens according to a binary masking indicator M E RR-. and replace them with the learnable mask tokens E RD before feeding them into the encoder. All the latent tokens are maintained to more heavily aggregate information on the unmasked image tokens and used to reconstruct the masked tokens at the decoder output.
3.1. Architecture
We build MAETok upon the recent ID tokenizer design with learnable latent tokens (Yu et al., 2024c; Li al., 2024c; Chen et al., 2024a). Both the encoder & and decoder D adopt the Vision Transformer (ViT) architecture (Dosovit- skiy et al., 2021; Yu et al., 2021), but are adapted to handle both image tokens and latent tokens, as shown in Fig. 3.
Encoder. The encoder first divides the input image I E x x3 into N patches according to predefined patch size P, each mapped to an embedding vector of dimension D, resulting in image tokens x E NDD, In addition, we define a set of L learnable latent tokens z E R. The encoder transformer takes the concatenation of image patch embeddings and latent tokens [x: E (N+ as its input, and outputs the latent representations E R with a dimension of H from only the latent tokens:
Auxiliary Shallow Decoders. In MAE, a shallow decoder (He al., 2022) or a linear layer (Xie et al., 2022; Wei et al., 2022) is required to predict the target features, e.g., raw pixel values, HOG features, and features from pre-trained models, of the masked image tokens from the remaining ones. However, since our goal is to train MAE as tokenizers, the pixel decoder 'D needs to be able to reconstruct images in high fidelity. Thus, we keep D as a similar capacity to E, and incorporate auxiliary shallow decoders to predict
E ([x:2) 2)).
(4)
4

---

Masked Autocncoders Are Effective Tokenizers for Difflusion Models
additional feature targets, which share the same design as the main pixel decoder but with fewer layers. Formally, each auxiliary decoder D aux takes the latent representations h
fine-tuning, we linearly decrease the mask ratio from 60% to 0% over 50K iterations, with the same training loss. More training details of tokenizers are shown in Appendix B.
and concatenate with their own d as imputs, and output y
Implementation Details of Diffusson Models. We use SiT (Li et al., 2024a) and LightningDiT (Yao & Wang, 2025) for diffusion-based image generation tasks after training MAETok. We set the patch size of them I and use a ID position embedding, and follow their original training setting for other parameters. We use SYT-L of 458M parame- ters for the analysis and ablation study. For main results, we train SiT-XL of 675M parameters for 4M steps and Light- Din for 400K steps on ImageNet of resolution 256 and 512. More details are provided in Appendix B.2.
as the reconstruction of their feature target y² E RNNDD
(7) where D' denotes the dimension of target features. We train these auxiliary decoders along with our AE using additional MSE losses at only the masked tokens according the masking indicator M, similarly to Xie et al. (2022):
y = ( ((e'; h):
maskk
Σ lI @ (5 - -
Evaluation. For tokenizer evaluation, we report the recon- struction Frechet Inception Distance (rFID) (Heusel et al., 2017), peak-signal-to- similarity index measure (SSIM) on ImageNet and MS- COCO (Lin et al., 2014) validation set. For the latent space evaluation of the tokenizer, we conduct linear probing (LP) on the flatten latent representations and report accuracy. To evaluate the performance of generation tasks, we report gen- eration FID (gFID), Inception Score (IS) (Salimans et al., 2016), Precision and Recall (Kynkäänniemi et al., 2019) (in Appendix C.1), with and without classifier-free guidance (CFG) (Ho & Salimans, 2022), using 250 inference steps.
3.3. Pixel Decoder Fine-Tuning
ratio (PSNR), and structural
While mask modeling encourages the encoder to learn a better latent space, high mask ratios can degrade immediate reconstruction. To address this, after training AEs with mask modeling, we freeze the encoder, thus preserving the latent representations, and fine-tune only the pixel decoder for small number of additional epochs. This process allows the decoder to adapt more closely to frozen latent codes of clean images, recovering the details lost during masked training. We use the same loss as in Eq. (6) for pixel decoder fine-tuning and discard all auxiliary decoders in this stage.
4.2. Design Choices of MAETok
4. Experiments
We first present an extensive ablation study to understand how mask modeling and different designs affect the recon- struction of tokenizer and, more importantly. the generation of diffusion models. We start with an AE and add different components to study both rFID of AE and gFID of SiT-L..
We conduct comprehensive experiments to validate the de- sign choices of MAETok, analyze its latent space, and bench- mark the generation performance to show its superiority.
4.1. Experiments Setup
Mask Modeling. In Table 1a, we compare AE and VAE with mask modeling and also study the proposed fine-tuning of the pixel decoder. For AE, mask modeling significantly improves gFID and slightly deteriorates rFID, which can be recovered through the decoder fine-tuning stage without sacrificing generation performance. In contrast, mask mod- eling only marginally improves the gFID of VAE, since the imposed KI. constraint may hinder latent space learning.
Implementation Details of Tokenizer. We use XQ-GAN codebase (Li et al., 2024d) to train MAETok. We use ViT- Base (Dosovitskiy et al., 2021), initialized from scratch, for both the encoder and the pixel decoder, which total have 176M parameters. We set L 128 and H = 32 for latent space. Three MAETok variants are trained on 256x256 ImageNet (Deng et al., 2009), and 512x512 ImageNet, and subset of 152 LAION-COCO (Schuhmann et al., 2022) for 500K iterations, respectively. In the first stage training with mask modeling on ImageNet, we adopt a mask ratio of 40-60%. set by ablation, and auxiliary shallow de- coders for multiple targets of HOG (Dalal & Triggs, 2005), DINO-v2-Large (Oquab et al., 2023), and Ligre (Zhai et al., 2023) features. We adopt additional auxil- iary decoder for tokenizer trained on LAION-COCO, which predicts the discrete indices of text captions for the image using& a BPE tokenizer (Cherti et al., 2023; Huang et al., 2024). Each auxiliary decoder has 3 layers also set by abla- tion. We set X = 1.0 and X2- = 0.4. For the pixel decoder
Reconstruction Target. In Table Ib, we study how different reconstruction targets affect latent space learning in mask modeling. We show that using the low-level reconstruction features, such as the raw pixel (with only a pixel decoder) and HOG features, can already learn a better latent space, re- sulting a lower gFID. Adopting semantic teachers such as DINO-v2 and CLIP instead can significantly improve gFID. Combining different reconstruction targets can achieve balance in reconstruction fidelity and generation quality.
Mask Ratio In Table 1c, we show the importance of proper mask ratio for learning the latent space using HOG target,
5

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
case
FFID
gFID
case
rFTD
gFID
low
high
FFID
gFID
blocks
rFID
gFlD
VAE
1.22
22.17
pixel
1.15
17.18
0
60
0.82
24.15
linear
1.35
6.98
+MM
1.75
18.17
HOG
2.43
13.54
10
40
1.01
22.63
1
1.19
6.43
AE
0.67
24.47
DINO
0.89
6.24
20
1.44
20.35
0.85
5.78
+MM
0.85
5.78
CLIP
0.78
11.31
40
40
1.78
18.27
6
0.86
7.12
0.48
5.69
Comb.
0.85
5.78
40
60
2.43
17.18
12
0.96
8.80
(a) Mask modeling.
(b) Reconstruction target.
(c) Mask ratio (HOG w/o FT).
(d) Aux. decoder depth.
Table Ablations with MAETok on 256x256 ImageNet. We report rFID of tokenizer and gFID of SiT-L trained on latent space of the
tokenizer without classifier-fiee guidance. We train tokenizer of 250K and SiT-L 400K steps. Default settings are indicated Grey
MART&
NAB
MARTAA CCUP.
KART&& (Pai
MAKTeR
60
REFA
13
16
AE
(b) VAE
MAETok
Figure 4. UMAP visualization on ImageNet of learned latent space from (a) AE: (b) VAE: (c) MAETok. Colors indicate differ- ent classes. MAETok presents more discriminative latent space.
30
600
20
10.3K
200
300K
00.00
Linesr Probing Aueussy
Toinning Steps
(a) gFID vs. LP Acc.
(b) gFDg during training
Figure 5. The latent space from tokenizer correlates strongly with generation performance, More discriminative latent space (a) with higher linear probing (LP) accuracy usually leads to better gFID, and (b) makes the learning of the diffusion model easier and faster.
highlighted in previous works (He et al., 2022; Wei al., 2022; Xie al., 2022). A low mask ratio prevents the AE from learning more discriminative latent space. A high mask ratio imposes a trade-off between reconstruction fidelity and the latent space quality, and thus generation performance.
tive and separated latent representation of MAETok results in much fewer GMM modes and improve the generation performance. More visualization is shown in Appendix C.3.
Auxiliary Decoder Depth. We study the depth of auxiliary decoder in Table 1d with multiple reconstruction targets. We show that decoder that is too shallow too deep could hurt both the reconstruction fidelity and generation quality. When the decoder is too shallow, the combined target fea- tures may confuse the latent with high-level semantics and low-level details, resulting in a worse reconstruction fidelity. However, deeper auxiliary decoder may learn a less dis- criminative latent space of the AE with its strong capacity, and thus also lead worse generation performance.
Latent Distribution and Generation Performance. We assess the latent space's quality by studying the relationship between the linear probing (LP) accuracy on the latent space, as a proxy of how well semantic information is preserved in the latent codes, and the gFID for generation performance. In Fig. 5a, we observe tokenizers with more discriminative latent distributions, as indicated by higher LP accuracy, cor- respondingly achieve lower gFID. This finding suggests that when features well-clustered in latent space, the genera- tor can more easily learn to generate high-fidelity samples. We further verify this intuition by tracking gFID) throughout training, shown in Fig. 5b, where MAETok enables faster convergence, with gFID rapidly decreasing with lower val- ues than the AE or VAE baselines. A high-quality latent distribution is shown be a crucial factor in both achieving strong final generation metrics and accelerating training.
We include more ablation study on the number of learnable latent tokens and 2D RoPE in Appendix C.4.
4.3. Latent Space Analysis
We further analyze the relationship between the latent space of the AE variants and the generation performance of SiT-L.
Latent Space Visualization. We provide a UMAP vi- intuitively
4.4. Main Results
sualization (MeInnes et al., 2018) in Fig. compare the latent space learned by different variants of AE, Notably, both the AE and VAE exhibit more entan- gled latent embeddings, where samples corresponding to different classes tend to overlap substantially. In contrast, MAETok shows distinctly separated clusters with relatively clear boundaries between classes, suggesting that MAETok learns more discriminative latent representations. In line with our analysis in Section 2 and Fig. 2, more discrimina-
Generation. We compare SiT-XL and LightningDiT based on variants of MAETok Tables and for the 256 and 5 512x512 ImageNet benchmarks, respectively, against other SOTA generative models. Notably, the naive SIT-XL trained on MAETok with only 128 tokens and plain AE architecture achieves consistently better gFID and IS with- out using CFG: it outperforms REPA (Yu et al., 2024d) by 3.59 gFID on 256 resolution and establishes a SOTA com-
6

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
w/o CFG
w/ CFG
Model (G)
Purams (G)
Model (T)
# Params (T)
Tokens L
rFID
gFID
IS
gRID J.
IS
AMto-regressvve
VQGAN (Esser 2021)
1.4B
VQ
23M
256
7.94
5.20
290.3
VIT-VQGAN (Yu et al., 2021)
1.7B
VQ
64M
1024
1.28
4.17
175.1
RQ-Trans. (Lee al., 2022)
3.8B
RQ
66M
256
3.20
3.80
323.7
MaskGrT (Chang al., 2022)
227M
VQ
66M
256
2.28
6.18
182.1
LlamaGen-3B (Sun et al., 2024)
3.1B
VQ
72M
576
2.19
2.18
263.3
TYTok-S-128 (Yu al., 2024c)
287M
VQ
72M
128
1.61
1.97
281.8
VAR (Tian al., 2024)
2B
MSRQ
109M
680
0.90
1.92
323.1
ImageHolder (1.i et al., 2024c)
362M
MSRQ
176M
286
0.80
2.60
293.0
MAGVTE-22 (Yu et al., 2024a)
307M
LFQ
116M
256
1.61
3.07
213.1
1.78
319.4
MaskBit (Weber et al., 2024)
305M
LFQ
54M
256
1.61
1.52
328.6
MAR-H (Li et ul., 2024b)
943M
KL
66M
256
1.22
2.35
227.8
1.55
303.7
baased LDM-4 (Rombach al., 2022h)
400M
KL
55M
4096
0.27
10.56
103.5
3.60
247.7
U-VIT-H/2 (Bao ul., 2023)
501M
2.29
263.9
MDTV2-X1/2 (Gao al., 2023)
676M
5.06
155.6
1.58
314.7
DIT-XL/2 (Peebles & Xic. 2023)
675M
KL
84M
1024
0.62
9.62
121.5
2.27
278.2
SIT-XL/2 (Ma et al., 2024)
8.30
131.7
2.06
270.3
REPA (Yu et al., 2024d)
5.90
157.8
1.42
305.7
TexTok (Zha et al., 2024)
675M
KL
176M
256
0.69
1.46
303.1
LightningDiT (Yao & Wang, 2025)
675M
KI.
70M
256
0.28
2.17
205.6
1.35
293.3
Ours
MAETоk LightinnDDi
675M
208.3
1.73
308.4
AE
176M
128
0.48
MAETok SiT-XL
675M
2.31
216.5
1.67
311.2
Table System-level comparison on ImageNet 256x conditional generation. SiT-XI. and LightningDii trained on MAETek achieves performance comparable state-of-the-ar using plain AE with only tokens. "Model (G)": the generation model. "# Params (G)": the number of generator's parameters. "Model (T)": the tokenizer model. "# Params (T)": the number of tokenizer's parameters. "# Tokens": the number latent tokens used during generation. indicates that the model has been trained on other data than ImageNet.
w/o
w/ CFG
Model (G)
Params (G)
Model (T) # Params (T)
Tokens
IFD
gFID
IS gFID J. IS
GAN
BigGAN (Chang et al., 2022)
8.43
177.9
StyleGAN-XL -L (Karras et al., 2019)
168M
2.41
267.7
Auto-regressive
MaskGIT (Chang et al., 2022)
227M
VQ
66M
1024
1.97
7.32
156.0
TTok---64 (Yu al., 2024c)
177M
VQ
202M
128
1.52
2.13
261.2
MAGVIT-v2 (Yu al., 2024a)
307M
LFQ
116M
1024
1.91
324.3
MAR-H H (1.i al., 2024b)
943M
KI.
66M
1024
2.74
205.2
1.73
279.9
Diffission-bssed
ADM (Dhariwal & Nichol, 2021)
23.24
58.06
3.85
U-VIT-H/4 (Bao al., 2023)
501M
4.05
263.8
DiT-XI /2 (Peebles Xie, 2023)
675M
KI.
84M
4096
0.62
9.62
121.5
3.04
240.8
SIT-XL/2 (Ma al., 2024)
675M
2.62
252.2
DIT-XL (Chen al., 2024b)
675M
9.56
2.84
UVTT-H (Chen ul., 2024b)
501M
9.83
2.53
UVIT-H (Chen al., 2024b)
501M
AE
323M
256
0.22
12.26
2.66
UVTT-28 (Chen et al., 2024h)
2B
6.50
2.25
US/T-2B (Chen et al., 2024b)
2B
2.90
1.72
Ours
MAETok + LightningDiT
675M
2.56
224.5
1.72
307.3
MAETok SIT-XL
675M
176M
128
0.62
2.79
204.3
1.69
304.2
MAET>k USiT-2B
1.72
244.3
1.65
312.5
Table System-level comparison on ImageNet 512 512 conditional gemeration. SiT-XL and LightningDiT trained on MAETok achieve
state-of-the-art performance using plain AE with only 128 tokens, outperforming USIT of 2B parameters using only 675M parameters.
parable gFID of 2.79 at 512 resolution. When using CFG, SiT-XL achieves a comparable performance with compet- ing autoregressive and diffusion-bar based baselines trained on VAEs at 256 resolution. beats the 2B USiT (Chen et al., 2024b) with 256 tokens and also achieves new SOTA of 1.69 gFID and 304.2 IS 512 resolution. Better results
have been observed with LightningDiT, trained with more advanced tricks (Yao & Wang, 2025), where it outperforms MAR-H of 1B parameters and USIT of 2B parameters with- out CFG, achieves 2.56 gFID and 224.5 IS, and 1.72 gFID with CFG. When using a USiT-2B (Chen et al., 2024b) for 512 generation, pushes the gFID without CFG to 1.72.

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
ImngeNet
COCO
from our results is that diffusion models trained MAETok usually present significantly better generation performance without CFG, compared to previous methods, yet smaller performance gap with CFG. We hypothesize that the reason is that the unconditional class also learns the semantics in
Tokenizer
Parans
Tokens
IFID
PSNR*
SSIM?
dFDD
PSNR
SSIM
256
SD-VAE!
84M
1024
0.62
26.04
0.834
4.07
25.76
0.845
DC-AET
323M
64
0.77
23.93
0.766
5.10
23.59
0.776
VA-NAE
70M
256
0.28
26.30
0.846
2.80
26.12
0.856
SoRVO
176M
64
0.61
22.97
0.739
5.16
22.86
0.745
ToxTok
176M
256
0.69
24.38
0.645
the latent space, as shown by the unconditional generation
MAETaK
176M
128
0.48
23.61
0.763
4.87
23.31
0.773
512x512
Metric AB
VAB
MAETek
MAETek
MAETek MAETek (DINO)
SD-VAE!
4096
0.19
27.36
0.849
2.41
26.48
0.841
(HOG)
(CIIP)
DC-AET
323M
256
0.21
26.23
0.815
2.85
25,47
0.811
gRID
59.02
58.34
45.31
34.73
20.76
18.31
TexTak
176M
24.45
IS
16.91
17.36
24.25
28.33
44.51
47.33
MAETak
176M
128
0.62
22.18
0.701
5.91
22.48
0.695
MAETak
176M
128
0.76
22.43
0.217
5.25
23.35
0.684
Table 5. Unconditional generation performance of SiT-L.
Table Comparison of various continuous tokenizers. indicates the tokenizer is trained other data than ImageNet. MAETok achieves a better trade-off of compression and reconstruction.
performance in Table 5. As the latent space becomes more discriminative, the unconditional generation performance also improves significantly. This implies that the CFG lin- ear combination scheme may become less effective (Zhao & Schwing, 2025), aligning with our CFG tuning results included in Appendix C.2. Moreover, adopting more recent advanced CFG techmiques, such as Autoguidance (Karras et al., 2024) with naively earlier checkpoints of the genera- tive model, and guidance- free training (Chen et al., 2025) can further improve the gFID of the SiT-XL model from 1.67 to 1.54 and 1.51, respectively. We leave the exploration of other CFG techniques for future work.
and gFIl) with CFG to 1.65. These results demonstrate that the structure of the latent space (see Fig. 4), instead of the variational form of tokenizers, is vital for the diffusion model to learn effectively and efficiently. We show a few selected generation samples in Fig. 1, and more uncurated visualizations are included in Appendix C.5.
Reconstruction. MAETok also offers strong reconstruc- tion capabilities on ImageNet and MS-COCO, as shown in Table 4. Compared to previous continuous tokenizers, including SD-VAE (Rombach et al., 2022a), DC-AE (Chen et al., 2024b), VA-VAE (Yao & Wang, 2025), SoftVQ-VAB (Chen et al., 2024a), and TexTok (Zha et al., 2024), MAE- Tok achieves a favorable trade-off between the quality of the reconstruction and the size of the latent space. On 256 256 ImageNet, using 128 tokens, MAETok attains an rFID of 0.48 and SSIM of 0.763, outperforming methods such as VQ in terms of both fidelity and perceptual similarity, while using half of the tokens in TexTok (Zha et al., 2024). On MS-COCO, where the tokenizer is not directly trained, MAETok still delivers robust reconstructions. At resolu- tion of 512, MAETok maintains its advantage by balancing compression ratio and the reconstruction quality.
Learnable Tokens, Mask Modeling, and Auxiliary De- coders. We also provide a study of the effect of each com- ponent in MAETok, as shown in Table 6. The results reveal clear and complementary gains from each design choice. When all three are present, MAETok reaches a better trade- off between the reconstruction and generation quality; drop- ping the auxiliary decoder or the masking objective, for instance, leads to noticeably weaker results, while omit- ting the learnable tokens, i.e., using only 256 image tokens, impairs both fidelity metrics most severely.
Mask
Learnable Token
Aux.
rFID
gFID
Modeling
Decoder
0.85
5.78
V
0.64
8.44
1.01
6.85
V
1.15
17.18
4.5. Discussion
v
0.43
9.88
0.96
18.23
Efficient Training and Generation. A prominent benefit of the ID tokenizer design is that it enables arbitrary number of latent tokens. The 256x256 and 512x512 images are usually encoded to 256 and 1024 tokens, while MAETok uses 128 tokens for both. It allows for much more efficient training and inference of diffusion models. For example, when using 1024 tokens of 512x5 112 images, the Gflops and the inference throughput of SIT-XL are 373.3 and 0.1 images/second on a single A 100, respectively. MAETok reduces the Glops to 48.5 and increases throughput to 3.12 images/second. With improved convergence, MAETok en- ables a 76x faster training to perform similarly to REPA.
0.67 24.47
Table 6. Effectts of different components in MAETok.
5. Related Work
Image Tokenization. Imgae tokenization aims at trans- forming the high-dimension images into more compact and structured latent representations. Farly explorations mainly used autoencoders (Hinton & Salakhutdinov, 2006; Vin- cent et al., 2008), which learn latent codes reduced dimen- sionality. These foundations soon inspired methods with variational posteriors, such as VAEs (Van Den Oord et al.,
Unconditional Generation. An interesting observation
8

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
2017; Razavi et al., 2019a) and VQ-GAN (Esser et al., 2021; Razavi et al., 2019b). Recent work has further improved compression fidelity and scalability (Lee et al., 2022; Yu et al., 2024a; Mentzer et al., 2023; Zhu et al., 2024), showing the importance of latent structure. More recent efforts have shown methods that bridge high- high-fidelity reconstruction and semantic understanding within a single tokenizer (Yu et al., 2024c; Li et al., 2024c; Chen et al., 2024a; Wu al., 2024; Gu et al., 2023). Complementary them, further high- light the importance of discriminative latent space, which allows us use simple AE yet achieve better generation.
Impact Statement
This work advances the fundamental understanding and tech- nical capabilities of machine learning systems, specifically in the domain of image generation through diffusion models. While contributions primarily technical, improving efficiency and effectiveness of generative models, we ac- knowledge that advances in image synthesis technology can have broader societal implications. These may include both beneficial applications in creative tools and design, as well as potential concerns regarding synthetic media. We have focused on developing more efficient and robust methods for image generation, and we encourage ongoing discussion about the responsible deployment of such technologies.
The paradigms of image generation mainly categorize to autoregressive and diffusion models. Autoregressive models initially relied on CNN architectures (Van den Oord al., 2016) and were later augmented with Transformer- bbasodedr---dddddddd models (Vaswani et al., 2023; Yu al., 2024b; Lee et al., 2022; Liu et al., 2024; Sun et al., 2024) for improved scalability (Chang et al., 2022; Tian et al., 2024). Diffusion models show strong performance since their debut Sohl-Dickstein et al. (2015b). Key developments (Nichol Dhariwal, 2021; Dhariwal Nichol, 2021; Song et al., 2022) refined the denoising process for sharper samples. A pivotal step in performance and efficiency came with latent diffusion (Vahdat et al., 2021; Rombach et al., 2022b), which uses tokenizers to reduce dimension and conduct denoising in a compact latent space (Van Den Oord et al., 2017; Esser et al., 2021; Peebles & Xie, 2023; Qiu et al., 2025). Recent advances include designing better tokenizers (Chen al., 2024a; Zha al., 2024; Yao & Wang, 2025) and combining diffusion with autoregressive models (Li et al., 2024b).
Image Generation.
Acknowledge
The authors would like to thank the anonymous reviewers and area chair for their helpful comments. This research project has benefited from the Microsoft Accelerating Foun- dation Models Research (AFMR) grant program. Difan Zou and Yujin Han acknowledge the support from NSFC 62306252, Hong Kong ECS award 27309624, Guangdong NSF 2024A 151501244, and the central fund from HKU.
6. Conclusion
We presented a theoretical and empirical analysis of latent space properties for diffusion models, demonstrating that fewer modes latent distributions enable more effective learning and better generation quality. Based on these in- sights, we developed MAETok, which achieves state-of-the- art performance through mask modeling without requiring variational constraints. Using only 128 tokens, our approach significantly improves both computational efficiency and generation quality on ImageNet. Our findings establish that more discriminative latent space, rather than variational constraints, is crucial for effective diffusion models, opening new directions for efficient generative modeling at scale.

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
References
DalaI, N. and Triggs, B. Histograms of oriented gradients for human detection. In 2005 IEEE computer society conference computer vision and pattern recognition leee, 2005.
Assran, M., Duval, Q., Misra, L,, Bojanowski, P., Vincent, P., Rabbat, M., LeCun, Y., and Ballas, N. Self-supervised learning from images with joint- eantddined predictive architecture. Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 15619- 15629, 2023.
(CVPR'0S), volume 1, pp.
Deng. C.. Zh, D., Li, K., Guan, S., and Fan, H. Causal diffusion transformers for generative modeling. arXiv 2024.
preprint arXiv:2412.
Bao, F., Nie, S., Xue, K., Cao, Y., Li, C., Su, H., and Zhu, J. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 22669-22679, 2023.
Deng, J., Dong, W., Socher, R., Li, L..-J,, Li, K., and Fei-Fei, Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 148-255. Icee, 2009.
Benton, J., Bortoli, V., Doucet, A., and Deligiannidis, G. Nearly d-linear convergence bounds for diffusion models via stochastic localization. 2024.
Dhariwal, P. and Nichol, Diffusion models beat gans on image synthesis. Advances neural information processing systems, 34:8780-8794, 2021.
Chang, H., Zhang, H., Jiang, I.,, Liu, C.,, and Freeman, W. T.
Diakonikolas, I., Kane, D. M., Pittas, T., and Zarifis, N. Sq lower bounds for learning mixtures of separated and bounded covariance gaussians. In The Thirty Sixth Annual Conference on Learning Theory, pp. 2319-2349, PMLR, 2023.
Maskgit: Masked generative image transformer, 2022. URL https::t//// /arxivv..vv.. org/
04200.
Chen, H., Lee, H., and Lu, Improved analysis score- based generative modeling: User-friendly bounds under minimal smoothness assumptions. In International Con- ference Machine Learning, pp. 4735-4763. PMLR, 2023.
Dosovitskiy, A. and Brox, T. Generating images with percep- tual similarity metrics based on deep networks. Advances in neural information processing systems, 29, 2016.
Chen, H., Wang. Z., Li, X., Sun, X., Chen, F., Liu, J., Wang. J., Raj, B., Liu, Z., and Barsoum, E. Softvq- vae: Efficient 1-dimensional continuous tokenizer. arXiv preprint arXiv:2412.1 22..:2442..1988. 2024a.
Dosovitskiy. A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer,
M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image worth 16x16 words: Transformers for image recognition at scale, 2021. URL https: / //arxivv.v org/ abs///as///// /2000. 11929.
Chen, H., Jiang, K., Zheng, K., Chen, J., Su, H., and Zhu, Visual generation without guidance. arXiv preprint arXiv:2501. 15420, 2025.
Esser, P., Rombach, R., and Ommer, B. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pp. 12873-12882 2021.
Chen, J., Cai, H., Chen, J., Xie, E., Yang, S., Tang, H., Li, M., Lu, Y., and Han, S. Deep compression autoen- coder for efficient high-resolution diffusion models. arXiv preprint arXiv:2410. 10733, 2024b.
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J.,
Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F., et al. Scaling rectified flow transformers for high nigh-resolution image synthesis. In Forty-first International Conference Machine Learning, 2024.
Chen, S., Chewi, S., Li, J., Li, Y.. Salim, A., and Zhang, A. R. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. arXiv preprint arXiv:2209. 1215, 2022.
Gao, S., Zhou, P., Cheng, M.-M., and Yan, S. Mdtv2: Masked diffusion transformer is a strong image synthe- sizer. arXiv preprint arXiv:2303. 14389, 2023.
Cherti, M., Beaumont, R., Wightman, R., Wortsman, M.. Ilharco, G., Gordon, C., Schuhmann, C., Schmidt, L., and Jitsev, J. Reproducible scaling laws for con- trastive eeguuage-image learning. In Proceedings of the IFFF/CVF Conference on Computer Vision and Pattern Recognition, pp. 2818-2829, 2023.
Gatmiry, K., Kelner, J., and Lee, H. Learning mixtures of gaussians using diffusion models. arXiv preprint arXiv:2404. 18869, 2024.
Chung. H., Kim. J., Park, G. Y., Nam, H., and Ye, J. C. Cfg++: Manifold-constrainee classifier free guidance for diffusion models. arXiv preprint arXiv:2406.08070, 2024.
Ghosh, D., Hajishirzi, H., and Schmidt, L. Geneval: An object-focused framework for evaluating text-to-image alignment. Advances in Neural Information Processing Systems, 36, 2024.
10

---

Masked Autoencoders Are Effective Tokenizers for Diffusson Models
Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial networks. Communications of the ACM, (3(11))139-14 144, 2020.
Kingma, D. P. Auto-encoding variational bayes. arXiv preprint arXivv 2013.
Kynkäänniemi, T., Karras, T., Laine. S., Lehtinen, J., and Aila, T. Improved precision and recall metric for assess- ing generative models. Advances in neural information processing systems, 32, 2019.
Gu, Y., Wang, X., Ge, Y., Shan, Y., Qie, X., and Shou, M. Z. Rethinking the objectives of vector-quantized tokenizers for image synthesis. 2023. URL https: //arxivvorrg abs/ 2212.2222...22..2 03185.
Kynkäänniemi, T., Aittala, M., Karras, T., Laine, S., Aila, T., and Lehtinen, J. Applying guidance in limited interval improves sample and distribution quality in diffusion models. arXiv preprint UrXiv:2204.0772: 2024.
He, K., Chen, X., Xie, S., Li, Y., Dollár, P., and Girshick, R. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 16000-16009, 2022.
Larsen, A. B. L.. Sønderby, S. K., Larochelle, H.,, and Winther, O. Autoencoding beyond pixels using learned similarity metric. In International conference on machine learning, pp. 1558-1566. PMLR, 2016.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in Neural Information Processing Systems, 30, 2017.
Lee, D., Kim, C., Kim, S., Cho, M., and Han, W.-S. Autore- gressive image generation using residual quantization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 11523-11532, 2022.
Higgins, L, Matthey, L., Pal, A., Burgess, C. P., Glorot, X., Botvinick, M. M., Mohamed, S., and Lerchner, A. beta- vae: Learning basic visual concepts with constrained variational framework. ICLR, 3, 2017.
Li, H., Yang. J., Wang, K., Qiu, X., Chou, Y., Li, X., and Li, G. Scalable autoregressive image generation with mamba. arXiv preprint arXiv:2408. 2245,, 2024a.
Hinton, G. E. and Salakhutdinov, R. R. Reducing the di- mensionality of data with neural networks. science, 313 (5786):504-500 2006.
Li, T., Chang, H., Mishra, S. K., Zhang, H.,, Katabi, D., and Krishnan, D. Mage: Masked generative encoder to unify representation learning and image synthesis, 2023. URL https://arxie org/ abs/ 2211 09117.
Ho, J. and Salimans, T. Classifier-fres free diffusion guidance. arXiv preprint arXiv:2207. 12598, 2022.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 33:6840 (6300066811110 2020.
Li, T. Tian, Y., Li, H., Deng, M., and He, Autoregres- sive image generation without vector quantization, 2024b. URL https: / aawaa.iiivv. http:://wwww.//oeee.e...oooo abs 2206. 118...1111301 11838.
Huang, Z., Ye, Q.. Kang, B., Feng. J., and Fan, H. Classi- fication done right for vision-language pre-training. In NeurlPS, 2024.
Li, X., Chen, H., Qiu, K., Kuen, J., Gu, J., Raj, B., and Lin, Z. Imagefolder: Autoregressive image generation with folded tokens. arXiv preprint arXiv:2410.01756 2024c.
Isola, P., Zhu, J.-Y., Zhou, T., and Efros, A. A. Image- to-image translation with conditional adversarial net- works, 2018. URL https://aat /aaaviv./riv..../ org/ abs/
Li, X., Qiu, K., Chen, H., Kuen, J., Gu, J., Wang. J., Lin, Z., and Raj, B. Xq-gan: An open-source image tokenization framework for autoregressive generation. arXiv preprint arXiv:2412.01765 2024d.
Johnson, J., Alahi, A., and Fei-Fei, L. Perceptual losses for real-time style transfer and super-resolution. In Com- puter Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11 11-14, 2016, Pro- ceedings, Part II 14, pp. 694-711. Springer, 2016.
Lin, T.-Y., Maire, M., Belongie, S., Hays, J., Perona, P., Ra- manan, D., Dollár, P., and Zitnick, C. L. Microsoft coco: Common objects in context. In Computer Vision-ECCY 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pp. 740- 755. Springer, 2014.
Karras, T., Laine, S., and Aila, T. A style-based generator architecture for generative adversarial networks. In Pro- ceedings of the JEEE/CVF conference on computer vision and pattern recognition, pp. 4401-4410, 2019.
Liu, W., Zhuo L.,, Xin, Y., Xia, S., Gao, P., and Yue, X. Customize your visual autoregressive recipe with set au- toregressive modeling, arXiv preprint arXiv:2410. 200.5111, 2024.
Karras, T., Aittala, M., Kynkäänniemi, T., Lehtinen, J., Aila,
T., and Laine. S. Guiding a diffusion model with a bad version of itself. arXiv preprint wxXiv:2406.02507. 2024.
Loshchilov, I. Decoupled weight decay regularization. arXiv preprint 1000000,, 2017.

---

Masked Autocncoders Are Effective Tokenizers for Diffusson Models
Ma, N., Goldstein, M., Albergo, M. S., Boffi, N. M., Vanden- Bijnden, E., and Xie, S. Sit: Exploring flow and diffusion- based generative models with scalable interpolant trans- formers, arXiv preprint arXiv:2401.08774 087440 2024.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with la- tent diffusion models, 2022b. httpt:://aaaa org/ /tl/abb//22111 /tbos/2112.0000 10752.
McInnes, L., Healy, J., and Melville, Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint arXiv: watp:::082.00000 2018.
Salimans, T., Goodfellow, I., Zaremba, W., Cheung, V., Radford, A., and Chen, X. Improved techniques for training gans. Advances in Neural Information Processing Systems, 29, 2016.
Mentzer, F., Minnen, D., Agustsson, E., and Tschannen, M. Finite scalar quantization: Vq-vae made simple, 2023.
Schuhmann, C.,, Beaumont, R., Vencu, R., Gordon, C., Wightman, R., Cherti, M., Coombes, T., Katta, A., Mullis, C., Wortsman, M., et al. Laion-5b: An open large-scale dataset for training next generation image-text models. Advances in Neural Information Processing Systems, 35:
Nichol, A. and Dhariwal, P. Improved denoising diffusion probabilistic models, 2021. URL https: /arxiv. org/abs/ 2102. 09672.
25278-25291 2022.
Oquab, M., Darcet, T., Moutakanni, T. Vo, H. V., Szafraniec, M., Khalidov, V., Fernandez, P., Haziza, D., Massa, F., El-Nouby, A., Howes, R., Huang, P.-Y,, Xu, H., Sharma, V., Li, S.-W., Galuba, W., Rabbat, M., As- sran, M., Ballas, N., Synnaeve, G., Misra, I., Jegou, H., Mairal, J., Labatut, P., Joulin, A., and Bojanowski, Dinov2: Learning robust visual features without supervi- sion, 2023.
Shah, K., Chen, S., and Klivans, A. Learning mixtures of gaussians using the ddpm objective. Advances in Neural Information Processing Systems, 36: 19636-196499 2023.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan,, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International conference on machine learning, pp. 2256-22662 2265. PMLR, 2015a.
Peebles, W. and Xie, S. Scalable difflusion models with trans- formers, 2023. URL attps://arxiv. org/ abs/ 2212. 0.748.
Sohl-Dickstein, J.. Weiss, A., Maheswaranathan, N. and Ganguli, S. Deep unsupervised learning using nonequilibrium thermodynamics, 2015b. URL https: / / //arxiv.orrg oww.oo////aooo. abs /1503 03585.
Qiu, K., Li, X., Kuen, J., Chen, H., Xu, X., Gu, Luo, Y., Raj. B., Lin, Z., and Savvides, M. Robust latent matters: Boosting image generation with sampling error synthesis. arXiv preprint arXiv:2503.08354 2025.
Song, J., Meng. C., and Ermon, Denoising diffusion implicit models, 2022. URL htt.s: https:://axxiii.ortt aarivvorv/.//////// abs abs/2C10.. 02502.
Qu, I.., Zhang, H., Liu, Y., Wang, X., Jiang, Y., Gao, Y., Ye. H., Du, D. K., Yuan, Z., and Wu, X. Tokenflow: Unified image tokenizer for multimodal understanding and generation. arXiv preprint arXiv:24/2.03066 2024.
Sun, P., Jiang. Y., Chen, S., Zhang, S., Peng, B., Luo, P., and Yuan, Z. Autoregressive model beats diffusion: Llama for scalable image generation. arXiv preprint a4Xiv:2406.063 2024.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp.
Tian, K., Jiang, Y., Yuan, Z., Peng. B., and Wang, L. Visual autoregressive modeling: Scalable image generation via next-scale prediction, 2024. URL pettps://arxit org/ abs / 2400 02905.
PMLR, 2021.
Razavi, A., Van den Oord, A., and Vinyals, O. Generating diverse high-fidelity images with vq-vae-2. Advances in neural information processing systems, 32, 2019a.
Tschannen, M., Eastwood, C., and Mentzer, F. Givt: Gen- erative
transformers. In European Conference on Computer Vision, pp. 292-309, Springer, 2025.
Razavi, A., van den Oord, A., and Vinyals, O. Generating diverse high-fidelity images with vae-2 2019b. URL https: / pettps://axxti //aaxivv../////////////////aatv/ wwtpxivvoorr///oooooooo abs/ 90666.... 6.066.10 00446.
Tseng, H.-Y., Jiang, L., Liu, C., Yang, M.-H., and Yang, W. Regularizing generative adversarial networks under lim- ited data, 2021. URL 2104. 03310.
abs//
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusson models. In Proceedings of the IEEE/CVF con- ference on computer vision and pattern recognition, pp. 10684-10695, 2022a.
Vahdat, A., Kreis, K., and Kautz, J. Score-based generative modeling in latent space, 2021. URL https: //arxiv. h0tgabb/21000 05931...1..........
12

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
Van den Oord, A., Kalchbrenner, N., Espeholt, L.,, Vinyals, O., Graves, A., et al. Conditional image generation with pixelenn decoders. Advances in neural information pro- cessing systems, 29, 2016.
Yu, Q. He, J., Deng, X., Shen, X., and Chen, L.-C. Ran- domized autoregressive visual generation. arXiv preprint arXiv:2411.00776, 2024b.
Yu, Q.. Weber, M., Deng, X., Shen, X., Cremers, D., and Chen, L.-C. An image is worth 32 tokens for reconstruc- tion and generation. arxiv:
Van Den Oord, A., Vinyals, O., et al. Neural discrete rep- resentation learning. Advances in neural information processing systems, 30, 2017.
2024c.
Yu, S., Kwak, S., Jang, H., Jeong, J., Huang, J., Shin, J., and Xie, S. Representation alignment for generation: Training diffusion transformers is easier than you think. arXiv preprint arXiv:24 06940, 2024d.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need, 2023. URL https: /arp:v////vvvv///a org/
abs/1705 03762.
Zha, K., Yu, L.,, Fathi, A., Ross, D. A., Schmid, C., Katabi, D., and Gu, X. Language-guicn guddeedguided image tokenization for generation. arXiv preprint arXiv:2412.0577 2024.
Vincent, P., Larochelle, H., Bengio, Y., and Manzagol, P.-A. Extracting and composing robust features with denoising autoencoders. In Proceedings of the 25th international conference on Machine learning, pp. 1096-1103, 2008
Zhai, X., Mustafa, B., Kolesnikov, A., and Beyer, L. Sig- moid loss for language image pre-training. In Proceed- ings of the IEEE/CVF International Conference on Com- puter Vision, pp. 11975-11986, 2023.
Weber, M., Yu, L., Yu, Q.. Deng. X., Shen, X., Cremers, D., and Chen, L.-C. Maskbit: Embedding frre image gen- cration via bit tokens. arXiv preprint 1rXiv:2409.1621,, 2024.
Zhang, Q., Wang, Y., and Wang, Y. How mask matters: To- wards theoretical understandings of masked autoencoders. Advances in Neural Information Processing Systems, 35: 27127-27133 2022.
Wei, C., Fan, H., Xie, S., Wu, C.-Y., Yuille, A., and Feicht- enhofer, C. Masked feature prediction for self-supervised visual pre-training. In Proceedings of the IEEE/CVF Con- ference on Computer Vision and Pattern Recognition, pp. 14668-14678 2022.
Zhang, R., Isola, P., Efros, A. A., Shechtman, E., and Wang, O. The unreasonable effectiveness of deep features as a perceptual metric, 2018. URL /ttp:://axxivivo.ooee abs/1801.. 03924.
Wu, Y. Zhang, 7.,, Chen, J., Tang, H., Li, D., Fang, Y., Zhu, L.,, Xie, E., Yin, H., Yi, L., et al. Vila-u: a unified foundation model integrating visual understanding and generation. arXiv preprint arXiv:2409.04.22 2024.
03924
Zhao, X. and Schwing, A. G. Studying classifier (-free) guid- ance from a classifier-centric perspective. arXiv preprint arXiv:2503. 10t.28803.1068,, 2025.
Xie, Z., Zhang, Z., Cao, Y., Lin, Y., Bao, J., Yao, Z., Dai, Q.. and Hu, H. Simmim: simple framework for masked image modeling. In Proceedings of the IEEE/CVF con- ference on computer vision and pallern recognition, pp. 9653-9663, 2022.
Zhu, L., Wei, F., Lu, Y., and Chen, D. Scaling the codebook size of vggan to 100,000 with a utilization rate of 99%. arXiv preprint arXiv:2406. 1837, 2024.
Yao, and Wang, X. Reconstruction vs. generation: Taming optimization dilemma latent diffusion models. arXiv preprint 2rxv:2501.004233 2025.
Yu, J., Li, X., Koh, J. Y., Zhang, H., Pang, R., Qin, J., Ku, A., Xu, Y., Baldridge, J., and Wu, Y. Vector-quantized image modeling with improved vggan. arXiv preprint 0rXhv22110.04677 2021.
Yu, L., Lezama. Gundavarapu, N. B., Versari, L., Sohn, K., Minnen, D., Cheng, Y., Gupta, A., Gu, X., Haupt- mann, A. G., Gong, B., Yang, M.-H., Essa, I., Ross, D. A., and Jiang, L. Language model beats diffusion tokenizer is key to visual generation. In The Twelfth International Conference on Learning Representations, 2024a. URL https: :ttttpp:://pppeeeee h/tp///peeee...//c//o/.e/p n 1 dd=gzqrANCF4g.
13

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
Figure 6. Additional selected samples from 512x512 SIT-XL model MAETok. We classifier-fres guidance scale of 2.0.
A. Theoretical Analysis
We begin the theoretical analysis by introducing the preliminaries of the problem and the necessary notation. Following the empirical analysis setting, we first consider the latent data distribution is the GMM with K equally weighted Gaussians:
Preliminary.
K
Σ N(H..I),
(9)
K
Following the the training objective (Shah et al., 2023), we consider the score matching loss of DDPM at timestep is
min EE sw (x, t)
Vx (x)
(10)
where sw (x, () is the denoising network and log (x( is the oracle score. Under the GMM assumption, the explicit solution of score function Vx log (x) can be written as
Vx
Σ i=1
a (x)mi, -x,
(11)
where the weighting parameter is
exp(- |x 12
(x),
PÍ, == p1 exp(-1) -)).
(12)
5= exp(ppppppp x 13/112/22)
Therefore, consider the denosing neural network with the following format, that is
80, (x) Σ (x) X,
(13)
where
(xx( = exp( |xx Mr 11222
P4, := #4 exp(-1).
(14)
== exp( -|[x - 22//11/22)
14

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
Figure 7. Additional selected samples from 256x256 diffusion models MAETok. We classifier- freet guidance scale of 2.0.
Assumptions. To ensure the denoising network approximates the score function with sufficient accuracy, we consider the following three common assumptions, which constrain the training process from the perspectives of data quality (separability), good initialization (warm start), and regularity (bounded mean of target distribution) (Chen et al., 2022; 2023; Benton et al., 2024).
Assumption A.1. (Separation Assumption in (Shah et al., 2023)) For mixture of K Gaussians given by Equation 9, for every pair of components i,j E (1, K} with i # j. we assume that the separation between their means
|l 14
C (lo(mmin(K, (K,
(15)
for sufficiently large absolute constant C' > Assumption A.2. (Initialization Assumption in (Shah et al., 2023)) For each component E (1,2...., K}, we have an initialization μe with the property that
||
(0)
1 S C² Vlog(min(KK. d))
(16)
for sufficiently small absolute constant C' 0.
Assumption A.3. The maximum mean norm of the GMM in GMM 9 bounded as: max; || S B.
Remark A.4. By Assumption A.3, we could derive the second movement bound of as
Excop [||x||]]= f po(x))lx||2dx d B²
(17)
Then, we can have the following analysis,
Step 1: From Modes to Training Loss. The main conclusion required for our proof is derived from the following theorem, which provides the estimation error || - 14, || for DDPM with gradient descent under (11 noise, assuming that Assumptions A. and A.2 are satisfied.
Theorem A.5. (Adopted from Theorem 16 in Shah et al. (2023)) Let be a mixture of Gaussians Eq. (9) with center parameters d* {μi "......,.. pxx Rid satisfying the separation A. I, and suppose we have estimates 0 for the centers
15

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
VAVAE
17.5
VAB
MAETOk
Maximum Mean Norm
15.0
12.5
10.0
5.0
2.5
0.0
10
Number of Components
Figure 8. We compare the maximum mean norm across different numbers of components and observe that AE, VAE, VAVAE, and our method MAETok exhibit similar maximum mean norms. This suggests that these latent spaces share comparable prior upper bound B, supporting the rationalc for primarily considering the number of modes, i.c., K in Theorem 2.1.
such that the warm initialization Assumption A.2 is satisfied. For any E >> Eu and noise scale t where
So 1/poly(d) 1 (0)(),
gradient descent on the DDPM objective at noise scale outputs Ō {it P.y.... j&) such that min; li, - u4ll S with high probability. DDPM runs for H H' iterations and uses 2 n' number of samples where
H' = Olog((=-1 log d)), n' = (K /E²")
Theorem A.5 indicates that, achieve the same estimation error c, a data distribution with more modes requires more training samples. Fig. 8 demonstrates that different latent spaces exhibit nearly identical mean norm upper bounds, thus justifying our focus on analyzing the number of modes K.
Given in Theorem A.5 and based on Assumption A.3, Eq. (10), Eq. (11), Eq. (13), we further have
B Bi]ls0,(x4 - Vx, log log
wa, (x (x) - (x,
K
K
< 2F Σ
i (x:) (x:) minill 2F
> (wx,(xx) (x,
19]
1=1
S 22 (c2 + B²)
(18)
The $ hides constant term 2 and 4.
Therefore, consider a step size hy I 7, we can have the learned score function sp, (x) satisfies
Ny (62 B²) T
Σ h Bllse, (xxx) V Vxig log Px (13
S
(19)
T
Step 2: From Training Loss Samlping Error. the practical sampling process, we adopt an early stopping strategy to improve the generation quality. Specifically, we consider the interval t [0, 0.8] during the reverse process. Then, the following conclusion holds:
Theorem A.6. (Theorem 2.2. in (Chen al., 2023)) There is universal constant such that the following hold, Suppose that Assumption A.3 and Eq. (19) hold and the step sizes satisfy the following for some quantities of opp--.. 02z
oix
h& <
2 <
k ==
=1,,
(20)
Th-1 Cd
16

---

Masked Autoencoders Are Effective Tokenizers for Diffusion Models
Define IT = :=<==> ki.. For T 2 2, v 2 the exponential integrator scheme (6) with early stopping results a distribution 45-6 such that
KL(psllêT- S (d + exp(-T)
+ Tej -
dl1
(21)
In particular, when using proper choices of lug, the quantity II can be as small as o(1). For instance, as shown in Chen et al. (2023), can be proved that = O(1/N2) when using exponentially decreasing stepsize.
Then combing Theorem A.5 and Theorem A.6, we finally have Theorem A.7. Training DDPM for 11 2 H1' iterations and uses n n' number of samples where
H's (loog(=- logd)) dl), n' = a B /=²)
Then, there is universal constant C such that the following hold. Suppose that Assumptions A.3 and Equation 19 hold and the step sizes satisfy
he fx-1
1
<
S7
h ==,,..,, N
(22)
Cd
For T 2 2,0 S 1. the exponential integrator scheme (6) with early stopping results in Define 11 == -EE A
distribution Q5-6 such that
KL(psllrr-o) (d B2)+x2)) Nr(22 + +B) + &II.
(23)
where p is the data distribution and 9 is the sampling distribution.
In Theorem 2.1, we establish a connection between the training process and the sampling process, using KL-divergence as a metric to quantify the distance between the true data distribution and the sampled generated data distribution. It should be noted that both KI. divergence and Wasserstein Distance serve as tools for measuring the similarity between distributions. Under the specific assumption that the data distributions are Gaussian, the Wasserstein Distance reduces FID (i.e., the metric used in paper). Theorem 2.1 demonstrates that achieving the same sampling training samples for data distributions with greater number of modes (K). Consequently, under the constraint of limited training samples, the quality of images generated from training data distributions with more modes (K) tends to be worse compared to those with fewer modes.
necessitates larger number of
B. Experiments Setup
B.1. Training Details of AEs
We present the training details of MAETok in Table 7.
B.2. Training Details of Diffusion Models
We present the training details of SiT-XL and LightningDiT in Tables and 9, which mainly follows their original setup.
B.3. Training Details of GMM Models
In Fig. 2, we train our own AE, KL- VAE, and MAETok under exactly the same settings and use the pre-trained VAVAE (Yao & Wang, 2025). The evaluation in Fig. 2 is performed with the same latent size and input dimensions. Specially, for GMM in Fig. 2a, we first represent the original latent size as (N, H1, C)), where N refers the training sample size, H refers to the number of tokens, and C' refers to the channel size. Following the typical GMM training, performed the following steps: (1) Latents flatten: The latent size becomes (N, H x C). (2) Dimensionality Reduction: To avoid the curse of dimensionality, we consider PCA and select a fixed dimension K that results in an explained variance greater than 90%. This step makes the latent dimension (N, K), ensuring that all latent spaces have consistent dimensions. (3) Normalization: To avoid numerical instability and feature scale differences, we further standardize the latent data. (4) Fitting: We the data using GMM and return the negative log- likeehod losses (NLL). We train the GMM on the entire Imagenet with a
17

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
Configuration
Value
image resolution
256x256, 512x512
enc/dec hidden dimension
768
enc/dec #heads
12
enc/dec #layers
12
enc/dec patch size
16
enc/dec positional embedding
2D RoPE (image), ID APE (latent)
optimizer
AdamW (Loshchilov, 2017)
base learning
le
weight decay
le-4
optimizer momentum
81, 32
0.9, 0.95
global batch size
512
learning rate schedule
cosine
warmup steps
10K
training steps
500K
augmentation
horizontal flip. center crop
discriminator
DINOv2-S
discriminator weight discriminator start
0.4 with adaptive weight
30K
discriminator LeCAM
0.001 (Tseng el al., 2021)
perceptual weight evaluation metric
1.0
FID-50k
Table 7. Training configuration of MAETok on 256x256 and 512x512 ImageNet.
batch size of 256 on single NVIDIA A8000. It should be noted that distributed training would further optimize the fitting time. The training time for GMM components of 50, 100, and 200 is roughly 3, 8, and 11 hours, respectively.
For SiT-L loss in Fig. 2b, we train SiT-L on the latent space these four tokenizers for 400K iterations, using an optimizer of AdamW, a constant learning rate of 1e-4, and no weight decay.
C. Experiments Results
C.1. More Quantitative Generation Results
We provide the additional precision and recall evaluation on 256x256 and 512x512 ImageNet benchmarks in Table 10 and Table 11, respectively.
C.2. Classifier-free Guidance Tuning Results
We provide our CFG scale tuning results in Table 12, where we found the gFID with CFG changes significantly even with small guidance scales. Applying CFG interval (Kynkäänniemi et al., 2024) to cutout the high timesteps with CFG can mitigate this issue. However, is still extremely difficult to tune the guidance scale.
We use a guidance scale of 1.9 and an interval of [0, 0.75) for 256 SiT-XI and guidance scale of 1.8 and interval of [0, 0.75] for 256x256 LightningDiT to report the main results. For 512x models, use guidance scale of 1.5 and an interval of [0, 0.7] for SiT-XL and guidance scale of 1.6 with an interval of [0, 0.65) for LightningDiTTs main results. Note that models may present even better results with more trained CFG tuning.
We attribute the difficulty of tuning CFG to the semantics learned by the unconditional class, as we discussed in Section 4.5. Such semantics makes the linear scheme of CFG less effective, as reflected by the sudden change with small guidance values. Adopting and designing more advanced CFG schemes (Chung et al., 2024; Karras et al., 2024) may also helpful with this problem, and is left as our future work.
C.3. Latent Space Visualization
More latent space visualization of MAETok variants is included in Fig. 9. MAETok in general learns more discriminative latent space with fewer GMM models with differente reconstruction targets.
18

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
Configuration
Value
image resolution
566x256,512x515
hidden dimension
1152
#heads
#layers
28
patch size
positional embedding
ID sinusoidal
optimizer
AdamW (Loshchilov, 2017)
base learning rate
le-4
weight decay
0.0
optimizer momentum
$1. 32 = 0.9, 0.999
global batch size learning rate schedule
256
constant
training steps augmentation
4M
horizontal flip, center crop
diffusion sampler
Fuler-Maruyama
diffusion steps
250
evaluation suite
ADM (Dhariwal & Nichol, 2021)
evaluation metric
FID-50k
Table 8. Training configuration of SiT-XL on 256 and 512x512 ImageNet.
(a) MAETok (Pixel)
(b) MAETok (HOG)
(c) MAETok (DINOv2)
(d) MAETok (CLIP)
Figure 9. UMAP visualization ImageNet of the leamned latent space from MAETok with pixel target; (b) MAETok with HOG
target; (c) MAETok with DINOv2 target; (d) MAETok with CLIP target. MAETok presents a more discriminative latent space.
C.4. More Ablation Results
We present the ablation study on latent tokens and 2D RoPE in Table 13. One can observe from Table 13a that using learnable latent tokens is more effective than using image tokens only, and 128 latent tokens is enough to achieve similar reconstruction and downstream generation performance, compared to 256 tokens. Furthermore, 2D RoPE helps to generalize better on different resolutions, when trained with mixed resolution images.
C.5. More Qualitative Generation Results
19

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
Configuration
Value
image resolution
256x256, 512x512
hidden dimension
1152
#heads
#layers
28
patch size
positional embedding
ID RoPE
optimizer
AdamW (Loshchilov, 2017)
base learning rate
2e
weight decay
0.0
optimizer momentum
fr 32 0.9, 0.95
global batch size Icarning rate schedule
1024
constant 400K
training steps
augmentation
horizontal flip. center crop
additional loss
cosine loss
diffusion sampler
Euler
diffusion steps
250
evaluation suite
ADM (Dhariwal & Nichol, 2021)
evaluation metric
FID-50K
Table 9. Training configuration of LightningDiT on 256x256 and 512x512 ImageNet.
Figure 10. Uncurated generation results of 256 MAETok STT-XL. We CFG of 3.0. Class label "Loggerhead" (33).
20

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
w/o CPG
w/ CRG
Model (G)
# Params
Model
Params
# Tokens
IFID
RFID
IS
Prec
Recall
AFID
1S
Prec ? Recall 1
Auto-negeessive
VQGAN (Esser 2021)
1,4B
VQ
23M
256
7.94
5.20
290.3
VIT-VQGAN (Yo et 2021)
1.78
VQ
64M
1024
1.28
4.17
175.1
RQ-Trans, (Lee et 2022)
RQ
256
3.20
323.7
MaskGIT (Chang
2022)
227M
VQ
66M
256
2.28
6.18
182.1
0.80
0.51
MAGE (Li et al., 2023)
439M
VQ
(N/A)
256
6.93
195.8
LlamaGen-3B
et al., 2024)
3,18
VQ
72M
576
2.19
2.18
263.3
0.80
0.58
TiTbk S- 128 (Ya
2024c)
287M
VQ
72M
128
1.61
281.8
VAR (Tian et al., 2014)
2B
MSRQ
109M
680
0.90
1.92
323.1
0.82
0.60
ImageFolder
2024c)
362M
MSRQ
176M
286
0.80
1.22
323.
0.75
0.63
(Yo MAGVIT-v2
2024a)
307M
LPQ
116M
256
1.61
3.07
213.1
1.78
319.4
MaskBit (Weber et al., 2024)
305M
LFQ
54M
256
1.61
1.52
328.6
MAR-H et al., 2024b)
943M
KL
66M
256
1.22
2.35
227.8
0.79
0.62
1.55
303.7
0.81
0.62
Diflusion-bused
(Rembach U-VIT-HI2 (Bao et al, 2023)
2022b)
400M
KL
55M
0.27
10.56
0.87
501M
2.29
263.9
0.82
0.57
MDTV--XL/2 (Gao et al., 2023)
676M
5.06
155.6
0.72
0.66
1.58
314.7
0.79
0.65
DIT-XL/ (Peebles Xic, 2023)
KL
9.62
278.2
STT-XL/2 (Ma et al., 2024)
8.30
131.7
0.68
0167
2,06
270.3
0.82
0.59
675M
REPA (Yu et al., 2024d)
5,90
157.8
0.70
0.69
1.42
305.7
0.80
0.65
TexTok-256 (Zha al., 2024)
675M
KL
176M
256
0.69
303.
0.79
0.64
LightinggDT (Yao & Wang, 2025)
675M
KLJ1
70M
256
0.28
2,17
205.6
1.35
295.3
Owr
MAETok + LightnngDDi
675M
2.21
208.3
0.79
0.62
1.73
308.4
0.80
0.63
STT-XL
AE
176M
128
0.48
310.6
0.63
Table 10. System-level comparison ImageNet x 256 conditional generation, now also reporting Precision and Recall under both CFG and no-CFG settings. "Model (G)": generation model. "# Params (G)": the number of generator parameters. "Model (T)": the tokenizer model. "# Params (T)": the number of tokenizer parameters. "# Tokens": the number of latent tokens used during generation. indicates that the model has also been trained on data beyond ImageNet.
w/o CFG
CFG
Model (G)
Params (G)
Model (T) Params (T)
Tokens
IFD
gFID
gFID
1S
Recall
GAN
BigGAN (Chang al,, 2022)
8.43
177.9
StyleGAN-XC (Karrus al., 2019)
168M
2.41
267.7
Anto-regresssvve
MaskGIT (Chung et 2022)
227M
VQ
66M
1024
1.97
7.32
156.0
TrTkk-B-64 2024c)
202M
1.52
2.13
261.2
MAGVIT-v2 (Yu ct 2024a)
307M
LFQ
116M
1024
1.91
324.3
MAR-H 2024b)
1.73
279.9
0.77
0.61
Diffscioo-bssee
ADM (Dhuriwal & Nichol, 2021)
23.24
58.06
3.85
221.7
0.84
0.53
U-VT-HH (Bao et al., 2023)
501M
4,05
263.8
0.84
0.48
DIT-XL/ (Peebles Xie, 2023)
675M
KL
SAM
4096
0.62
9.62
121.5
3.04
240.8
0.84
0.54
SIT-X1/2 (Ma et al., 2024)
675M
2.62
252.2
0.84
0.57
DIT-XI. (Chen al. 2024h)
675M
9.56
2.84
UVIT- (Chen et al., 2024b)
501M
9.83
2.53
UWT-H (Chen
2024b)
2,66
UVTT-28 (Chea al. 2024b)
2B
6.50
2.25
USIT-2B (Chea et al., 20244b
2.90
1.72
MAETok- LightningDiT
675M
2.56
224.5
1.72
307.3
0.81
0.62
176M
MALTok SIT-XL
675M
2.79
204.3
0.81
0.62
1.69
304.2
0.82
0.62
Table System-level comparison on ImageNet 512x512 conditional generation, now also reporting Precision and Recall for both CFG and no-CFG settings. "Model (G)": generation model. "# Params (G)": number of generator parameters. "Model (T)": the tokenizer model. "# Params (T)": number tokenizer parameters. "# Tokens": number latent tokens used during generation. indicates the model was also trained on data beyond ImageNet.
CFG
1.7
1.8
1.9
2.0
1.8
1.9
1.9
2.0
1.7
1.8
Interval
[0. 0.7]
[0. 0.7]
[0. 0.7]
[0, 0.7]
[0. 0.75]
[0, 0.75]
[0, 0.755
[0. 0.8]
[0, 0.8]
1.0]
[0. 1.0]
[0.125. 0.8]
gFID
4.96
4.94
4.91
4.92
4.92
4.92
4.94
5.09
5.14
5.21
8.55
6.08
IS
267.87
275.87
282.52
288.78
290.47
299.36
306.31
318.41
326.97
304.58
349.97
289.27
Table 12, CFG tuning results of 256 SIT-XL trained for 2M steps. We compute gFID and IS using 10K generated samples.
21

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
Tok
ID
# Tokens
rFD
gFID
VAVAE
256
0.28
13.65
MAETok V
256
0.37
5.05
MAETok
256
1.01
6,85
MAETok
128
0.48
5.69
(a) Latent tokens.
Pos.
256 rFID 512 rFID
Emb.
APE
0.73
1.43
RoPE
0.51
0.72
RoPE.
Table 13. Ablations of latent tokens and 2D RoPE with MAETok on 256x256 ImageNet. report rFID of tokenizer and gFID of SiT-L trained latent space the tokenizer without classificr-free guidanee. We train tokenizer 250K and SiT-L for 400K steps.
Figure 11. Uncurated generation results of 256 MAETok + STT-XL.. We use CFG of 3.0. Class label = "Macaw" (88).
Figure 12. Uncurated generation results of 256 MAETok STT-XL. We CFG of 3.0. Class label "Cacatua galerita" (89).
22

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
Figure 13. Uncurated generation results of 256 MAETok + SIT-XL. We use CFG of 3.0. Class label = "Flamingo" (130).
Figure 14, Uncurated generation results of 512x512 MAETok SIT-XL, We use CFG of 2.0. Class label = "Siberian husky" (250).
23

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
Figure 15. Uncurated generation results of 512x512 MAETok + SIT-XL. We use CFG of 2.0. Class label = "Aretic fox" (279).
Figure 16. Uncurated generation results of 512 MAETok SIT-XL, We use CFG of 2.0. Class label = "Otter" (360).
24

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
Figure/ 17. Uncurated generation results of 512 MAETok + SIT-XL. We use CFG of 2.0. Class label = "Guitar" (402).
Figure 18. Uncurated generation results of 512 MAETok SIT-XL. We use CFG of 2.0. Class label = "Fire Truck" (555).
25

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
Figure 19. Uncurated generation results of 512x512 MAETok + SIT-XL. We use CFG of 2.0. Class label = "Go-kart" (573).
Figure 20. Uncurated generation results of 512 MAETok SIT-XL. We use CFG of 2.0. Class label = "Laptop" (620).
26

---

Masked Autocncoders Are Effective Tokenizers for Diffusion Models
Figure 21. Uncurated generation results of 512x512 MAETok + SIT-XL. We use CFG of 2.0. Class label = "Carriage" (705))
Figure 22. Uncurated generation results of 512 MAETok SIT-XL. We use CFG of 2.0. Class label = "Sports Car" (402).
27


Reviews:
[{"review_id": "5TKC4ga1Ry", "reviewer": "Reviewer_pSm9", "summary": "This paper studies the properties of latent space for diffusion models, and claims that a more discriminative latent space (fewer Gaussian mixture modes) enable more effective diffusion learning and generation quality. Specifically:\n1. The paper conducts both empirical and theoretical analysis to show that the latent space with fewer gaussian modes, and thus better separability achieves lower diffusion loss in training.\n2. Based on 1, the paper further proposes MAETok, a mask autoencoder architecture which facilitates discriminative feature learning, with the additional learning objectives of predicting semantically rich features like HOG, DinoV2, and CLIP features.\n3. Extensive experimental results show that the proposed MAETok achieves state of the art generation results with fewer latent tokens, and thus less computation load, demonstrating the efficacy and validity of the proposed idea.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your time and efforts reviewing our paper. We now address the raised questions as follows.\n\n----\n\n> Why would it be more difficult to apply linear CFG scheme on semantically rich latent features?\n\nThanks for this interesting question. \n\n\nWe believe the limitation here is to apply the linear CFG scheme between a **semantic unconditional model** and a **semantic conditional model**.\nThere could be two possible reasons why it is more difficult to apply CFG in this scenario. \nThe first the linear interpolation could be more often lie outside the latent distribution. where the tokenizer decoder cannot reconstruct/decode reasonably (Fig. 4).\nAnd the second is the guidance scale becomes more difficult to tune with an unconditional model that already learns a good/semantic distribution (Tab. 5) [1]. In this case, applying CFG to obtain a 'lower-temperature' distribution becomes harder since the original assumption of CFG requires the unconditional model to be free of condition (Eq. (30) in [2]’s Appendix H): $ \\hat{q}  ( x_{t+1}  | x_t , c)   \\triangleq q(x_{t+1} | x_t )$.\nHowever, this assumption does not always hold [3], especially in our case. \n\n\n> If learning a discriminative, semantically rich latent feature space means that it will be more difficult to adapt with CFG scheme, how to justify and value and necessity of MAETok under this setting? Since currently CFG with vanilla VAE features still achieve better results.\n\n\nThanks for this great question. \n\nTo justify the value and necessity of MAETok, we need to highlight:\n* **MAETok achieves better results than vanilla CFG with vanilla VAE**. MAETok utlizes an pure autoencoder architecture with **only 128 tokens**. It achieves comparable performance as previous models on 256x256 generation and better results on 512x512 generation with a **vanilla SiT-XL model**, where previous best results are achieved with more training techniques on diffusion models (autoregressive, masking, noise scheduling, etc.). MAETok can also be used with these diffusion training techniques to obtain even better performance.\n* **MAETok works better with more advanced CFG schemes**. The difficulty of adapting naive CFG with MAETok lies in the problem of vanilla CFG, as we discussed before. CFG itself is very tricky and need to be tuned to find the optimal guidance scale. Adopting more recent and advanced CFG schemes can help MAETok achieve better results, as shown below.\n\n| CFG Scheme \t| FID \t| IS \t|\n|:---:|:---:|:---:|\n| Vanilla \t| 1.73 \t| 308.4 \t|\n| + Bad Version [1] \t| 1.54 \t| 315.9 \t|\n| + GFT [4] \t| 1.51 \t|  312.5\t|\n\nThese results are obtained using MAETok with SiT-XL on 256 generation. For the bad version guidance [1], we utilized the same SiT-XL model but from 1/4 training (1M steps). Designing more proper bad version models could further improve the results as in [1]. \n\nWe will add these results in our revised paper.\n\n[1] Karras et al. Guiding a Diffusion Model with a Bad Version of Itself.\n[2] Dhariwal et al. Diffusion Models Beat GANs on Image Synthesis.\n[3] Zhao et al. Studying Classifier(-Free) Guidance From a Classifier-Centric Perspective.\n[4] Chen et al. Visual Generation Without Guidance. \n\n---- \n\nWe hope the above response could resolve your questions and if there is further concern please let us know.", "claims_and_evidence": "The claim that discriminative latent space enables effective generative learning is supported by clear and convincing empirical and theoretical analysis.\n1. The paper shows that among the different latent spaces learnt by different variations of auto encoders, the ones with less gaussian mixture modes, and thus better separability and discriminative achieve lower diffusion training loss\n2. The paper further shows theoretically that with all other conditions kept the same, the latent space that have more gaussian mixture modes require more data samples to achieve the same diffusion loss, further supporting that discriminative latent space facilitates effective diffusion learning with less data samples.\n3. Experimental results show the consistency between high linear probing accuracy in latent space (better discriminative features) and better generative quality (through gFID)", "methods_and_evaluation_criteria": "Yes,  the proposed MAETok adopts a set of training designs to facilitate discriminative feature learning (mask auto encoding, predicting HOG, CLIP, DINOV2); For evaluation, the paper uses linear probing in the latent space to measure discriminability in latent feature space, and also show the corresponding generative performance through gFID, demonstrating the better feature space discriminability indeed facilitate better generation quality.", "theoretical_claims": "I have looked through the proofs about theorem 2.1 in the supplementary, but I don't think I have strong enough theoretical background to identify any issues.", "experimental_designs_or_analyses": "Yes, I checked all the subsections of the Experiment section. Specifically, in section 4.5, the paper points out that with CFG, the proposed method shows worse gFID comparing to previous methods, and hypothesize that the reason is that MAETok already learns a semantically rich latent space, and that the linear scheme of CFG on top of that may not be effective, which is also backed by the tuning results in Appendix C.2. I have a question based on this observation: if learning a discriminative latent space hurts its adaptability with CFG, then how to justify the necessity of MAETok?", "relation_to_broader_scientific_literature": "To facilitate discriminative latent feature learning, the paper adopts prior ideas like mask autoencoding [1], and further use auxiliary decoders to reconstruct targets such as HOG feature [2], DINOv2 feature [3], and CLIP feature [4], which are known to have good discriminatability. \n\n[1] He, Kaiming, et al. \"Masked autoencoders are scalable vision learners.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.\n[2] Dalal, Navneet, and Bill Triggs. \"Histograms of oriented gradients for human detection.\" 2005 IEEE computer society conference on computer vision and pattern recognition (CVPR'05). Vol. 1. Ieee, 2005.\n[3] Oquab, Maxime, et al. \"Dinov2: Learning robust visual features without supervision.\" arXiv preprint arXiv:2304.07193 (2023).\n[4] Radford, Alec, et al. \"Learning transferable visual models from natural language supervision.\" International conference on machine learning. PmLR, 2021.", "essential_references_not_discussed": "I haven't found missing essential references.", "other_strengths_and_weaknesses": "1. The paper is well written and easy to follow\n2. Empirical and theoretical evidence are comprehensive and strongly support the claims", "other_comments_or_suggestions": "please see questions.", "questions_for_authors": "1. As mentioned in Experimental design and analysis section, based on the results shown at Table 2, CFG without MAETok shows better gFID than CFG with MAETok. I have two questions regarding this observation and the analysis in the paper.\n(1) Why would it be more difficult to apply linear CFG scheme on semantically rich latent features? \n(2) If learning a discriminative, semantically rich latent feature space means that it will be more difficult to adapt with CFG scheme, how to justify and value and necessity of MAETok under this setting? Since currently CFG with vanilla VAE features still achieve better results.", "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": "iryXCQKr0w", "reviewer": "Reviewer_zfeD", "summary": "In this work, the authors find that the latent space with fewer modes and more discriminative features are better for training latent diffusion models. They propose a masked autoencoder for learning the latent space, where the decoder is later finetuned with encoder frozen, and achieve state-of-the-art FID.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your suggestions on reviewing our paper. We now address the questions raised as follows.\n\n--- \n\n\n> While I wonder that, a more discriminative latent leads a lower FID, would this be related to how FID is computed? If a latent contains more discriminative information, it may also be reflected in the decoded pixel space, which is then used to compute the FID -- on the features of a deep neural network, where discriminative images may have more stable features.\n\n\nThanks for your question. A lower FID is largely independent of stable features, as FID reflects the comparison between the mean and covariance of the real and generated data distributions. On the other hand, more discriminative latents describe whether the target latent distributions are separable and are largely independent of the comparison of their mean and variance.\n\n\n\n> Besides FID, is faster convergence / better quality also observed from visual quality and human evaluation?\n\nYes. Faster convergence can be observed obviously with better visual quality for the images generated during training. \nWe list some visualization examples during the training of SiT-L on VAVAE and MAETok at this [Figure link](https://anonymous.4open.science/api/repo/Rebuttal-8656/file/ID14576_maetok_vavae.pdf?v=b6ade62e).\nThese visualizations are generated using 30 inference steps and a guidance scale of 1.5 during training.\n\nAlso we conducted a human evaluation on 100 images generated by SiT-L trained on VAVAE and MAETok. The win, tie, lose rate of MAETok over VAVAE is shown below:\n|  \t| Win \t| Tie \t| Lose \t|\n|:---:|:---:|---|---|\n| MAETok over VAVAE\t| 0.75 \t| 0.18 \t| 0.11 \t|\n\nWe will add the visualization and human evaluation results into our revised Appendix.\n\n> Typo\n\nThanks for pointing this out. We will fix this typo in our revised paper.\n\n\n> L134 \"the generation quality of diffusion models is dominated by the denoising network’s training loss\", does it mean smaller diffusion loss indicates better quality? Based on score matching, I think this may not be always true, especially for different latent spaces?\n\nSorry for the confusion caused. We found this sentence is indeed not very accurate. What we are trying to convey is that the quality of the learned distribution of diffusion models is dominated by its training loss. \nThe eventual generation quality will determined by both the quality of the learned distribution (from the tokenizer encoder) and the capacity of the tokenizer decoder. \nIn our theoretical analysis, we assume the tokenizer decoders present similar capacity, and in that case, the generation quality will be determined by the learned distribution and thus diffusion loss. We will fix this sentence in our revised paper to reduce confusion.\n\n\n--- \n\nWe hope the above response resolved the questions.", "claims_and_evidence": "Claims are well supported. Both the generation results and visualization of the features support the claims.", "methods_and_evaluation_criteria": "The method is consistent with the hypothesis. Evaluation metric is standard.", "theoretical_claims": "Theoretical proofs look good to me.", "experimental_designs_or_analyses": "Experiments are solid. The authors evaluate on standard imagenet benchmark and have ablation studies on different design choices.", "relation_to_broader_scientific_literature": "This work aims at answering the question in the literature -- \"what is a good latent space for training latent diffusion model\", with a masked autoencoder as the proposed method and state-of-the-art performance.", "essential_references_not_discussed": "I did not find missing key related works.", "other_strengths_and_weaknesses": "The findings are interesting and experiments are solid.\n\nThis may not be a major weakness, while I wonder that, a more discriminative latent leads a lower FID, would this be related to how FID is computed? If a latent contains more discriminative information, it may also be reflected in the decoded pixel space, which is then used to compuate the FID -- on the features of a deep neural network, where discriminative images may have more stable features. Besides FID, is faster convergence / better quality also observed from visual quality and human evaluation?", "other_comments_or_suggestions": "Typo L78: as as", "questions_for_authors": "L134 \"the generation quality of diffusion models is dominated by the denoising network’s training loss\", does it mean smaller diffusion loss indicates better quality? Based on score matching, I think this may not be always true, especially for different latent spaces?", "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": "gBu6mMxZVq", "reviewer": "Reviewer_t7ec", "summary": "This paper analyzes how to develop a good image tokenizer. The authors bridge the GMM model and the quality of the latent space for generation and provide interesting discussion. Based on their investigation, they introduce MAETok to regularize the latent space with mask modeling when training tokenizer and achieve promising generation results.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your time and efforts reviewing our paper. We now address the raised questions as follows.\n\n\n--- \n\n> Does the evaluation in Figure 2 perform at the same latent size? The authors should provide more detailed experimental settings about this experiment.\n\nSorry for the confusion. In Figure 2, we train our own AE, KL-VAE, and MAETok under exactly the same settings and use the pre-trained VAVAE. The evaluation in Figure 2 is performed with the same latent size and input dimensions. Specially,\n\n(1) For GMM in **Figure 2(a)**\n\nWe first represent the original latent size as $(N, H, C)$, where $N$ refers to the training sample size, $H$ refers to the number of tokens, and $C$ refers to the channel size. Following the typical GMM training, we performed the following steps:\n- **Latents flatten**: The latent size becomes $(N, H \\times C)$.\n- **Dimensionality Reduction**: To avoid the curse of dimensionality, we consider PCA and select a fixed dimension $K$ that results in an explained variance greater than 90%. This step makes the latent dimension $(N,K)$, ensuring that all latent spaces have consistent dimensions.\n- **Normalization**: To avoid numerical instability and feature scale differences, we further standardize the latent data.\n- **Fitting**: We fit the data using GMM and return the negative log-likelihood losses (NLL).\n\n\n(2) For SiT-L loss in **Figure 2(b)**\n- We train SiT-L on the latent space of these four tokenizers for 400K iterations, using an optimizer of AdamW, a constant learning rate of 1e-4, and no weight decay. \n\n\nWe will include these experiment details in our revised Appendix.\n\n\n> Can the finding in Figure 2 be a direct criterion to find a good tokenizer? How much time/computation resource?\n\n\nThanks for this good question. In our image generation setting, if the decoders of different tokenizers present similar capability (rFID), the finding in Figure 2 is can be used to verify on good encoder/latent space for downstream generative models. For broader tokenizers, such as those for multimodal or video data, whether the findings in Figure 2 are applicable requires further discussion. We will leave additional exploration of this for future work.\n\n\nThe time and computation for the GMM analysis of the latent space is cheap. Using the Tokenizer AE, we train the GMM on the entire Imagenet with a batch size of 256 on a single NVIDIA A8000 GPU. It should be noted that distributed training would further optimize the fitting time. We consider various components and report corresponding times as follows:\n\n\n| #Components| Time (h) |\n| -------- | -------- | \n| 50  |   3  | \n| 100  |  8   | \n| 200  |   11   | \n\nThe GMM analysis time is much less compared to the downstream generative model training time. \nFor example, training SiT-XL for 4M steps on 8xH100 GPUs takes at least one week.\n\n\n> Could you provide a comparison of semantic regularization tokenizer versus generative model (e.g., REPA)?\n\n\nThanks for this interesting question. \n- First, a gFID convergence comparison of semantic regularization tokenizers, e.g. MAETok, with REPA is already included in Figure 5b. We also have the system-level comparison of training plain SiT-XL on MAETok, with REPA, MDTv2 in Table 2 and Table 4, which are semantic regularization methods added into training of generative models. These results show that **adding semantic information to the tokenizer outperforms adding semantic information to the generative model, as evidenced by the smaller gFID**. This comparison demonstrates a more fundamental improvement from the latent space than generative models. \n- Secondly, **semantic regularization in the generative model and in the tokenizer are not mutually exclusive.** We added additional experiments with REPA and MAETok and observed that the gFID of training SiT-L on MAETok decreased from 5.69 to 5.12 with REPA. \n\n> Could you also provide an aligned latent size for better comparison?\n\nThanks for this great suggestion. We were using 128 tokens mainly for the efficiency consideration, especially for 512x512 generation with SiT-XL.\nWe would like to provide a comparison here on the aligned latent size of MAETok with 256 tokens for 256x256 generation with SiT-L in the following:\n\n| Tokens \t| LP \t| rFID \t| gFID \t|\n|:---:|:---:|:---:|:---:|\n| 128 \t|  72.3\t| 0.48 \t|  5.69\t|\n| 256 \t|  74.5\t|  0.37\t|  5.05\t|\n\nNote that we simply set the learnable token length as 256 for the aligned latent size. \nAblation results on using only the 256 image tokens without the learnable tokens can be found in our response to Reviewer cxaR.\nWe will add these results in our revised paper.\n\n> The authors should discuss with 'Rethinking the Objectives of Vector-Quantized Tokenizers for Image Synthesis'.\n\n\nThank you for bringing up this relevant paper, and we will include it in our reference with proper discussion in our revised version.\n\n--- \n\nThanks for the great suggestions and questions again. \nPlease let us know if there is further concern.", "claims_and_evidence": "This paper should provide more detailed experiments to support its claims.\n\n1) Does the evaluation in Figure 2 perform at the same latent size? The authors should provide more detailed experimental settings about this experiment.\n2) Can the finding in Figure 2 be a direct criterion to find a good tokenizer? How much time/computation resource is needed to fit the GMM model, compared to directly training a downstream generative model?\n3) Could you provide a comparison of semantic regularization added in training the tokenizer versus training the generative model (e.g., REPA)?\n4) Could you also provide an aligned latent size (e.g., 128 tokens, 256 tokens) for better comparison with previous works, since your contribution of MAE is not directly related to latent size?", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "I have checked the correctness of theoretical claims.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "This paper relates to a broader effort to design tokenizers and find good latent spaces for generative models.", "essential_references_not_discussed": "The authors should discuss with 'Rethinking the Objectives of Vector-Quantized Tokenizers for Image Synthesis,' which examines the dillema between reconstruction and generation in vector-quantized tokenizers.", "other_strengths_and_weaknesses": "This paper reveals an interesting finding about the relationship between GMM mode and downstream generation performance. However, it has weaknesses in experimental design. See the \"Claims and Evidence\" part for my detailed comments on adding experiments to support the claims.", "other_comments_or_suggestions": "N/A.", "questions_for_authors": "N/A.", "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": "uioM837mcj", "reviewer": "Reviewer_cxaR", "summary": "This paper proposes to use masked autoencoder for reconstruction and verify it works better for diffusion model generation compared to AE and VAE.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thanks for your time reviewing this paper and suggestions on additional ablation studies and comparison results.\n\n \n---\n\n> It's a bit contradictive between figure 2 and 4.\n\nThanks for your question. \n\n- **Figure 4 and Figure 2 are aligned**. Figure 4 shows that the latent space of AE is more concentrated compared to others. In a concentrated latent space, increasing the number of modes in the GMM may not significantly improve fitting performance (e.g., reduce NLL) due to the close distances between modes, as observed in Figure 2.\n-  **Compactness in two dimensions (UMAP) is not highly correlated to its discriminative and modes**. To support our claim, we conduct a [synthetic experiment](https://anonymous.4open.science/api/repo/Rebuttal-8656/file/ID14576_%20Synthetic.pdf?v=def2b7f2) with two Gaussian Mixtures: A (30 close modes) and B (4 distant modes). UMAP shows that data with fewer modes (B) is more spread out, while data with more modes (A) is more compact.\n\n> Motivation of learnable token z? \n\nThe motivation of learnable token is for (1) more flexibility compression of images, and thus for more efficient downstream generative model training as discussed Sec. 4.5, and (2) for more compression purpose of MAE learning. \nThe following is an additional comparison on learnable tokens of the tokenizer and downstream SiT-L performance.\n\n| Latent  | # Tokens  | rFID  | gFID  |\n|:---:|:---:|:---:|:---:|\n| Learnable Token  | 128  |  0.85 | 5.78  |\n| Image Token  | 256  | 1.01  | 6.85  |\n\nWe will add this comparison in our revised paper. \n\n>  Ablation on 1. masked autoencoder, 2. learnable token; 3. auxiliary decoders.\n\nWe have provided ablation study mainly focusing on 1 (masked autoencoder) and 3 (auxiliary decoders) in Table 1. \nHere, we additional provide some ablation results on 2 (learnable tokens):\n\n| 1. masked autoencoder  | 2. learnable tokens  | 3. auxiliary decoders  | rFID  | gFID  |\n|:---:|:---:|:---:|:---:|:---:|\n| 1  | 1  | 1  | 0.85  | 5.78  |\n| 0  | 1  | 1  | 0.64  | 8.44  |\n| 1  | 0  | 1  | 1.01  | 6.85  |\n| 1  | 1  | 0  | 1.15  | 17.18  |\n| 0  | 0  | 1  | 0.43  | 9.88 |\n| 1  | 0  | 0  | 0.96    | 18.23 |   |\n| 0  | 1  | 0  | 0.67  | 24.47  |\n\nFrom the results, removing learnable tokens may lead to both decrease of rFID and gFID. \nMasking modeling is necessary for learning better latent space other than using only auxiliary decoders. \nWe will include these ablation results in our revised paper.\n\n> Elaborate details on figure 2.\n\nSorry for the confusion. We train AE, KL-VAE, and MAETok under the same settings and use the pre-trained VAVAE. The analysis is performed with the same latent size. Specifically:\n- **Latents Flatten** and **Dimensionality Reduction**: The flatten operation first changes the latent size from $(N, H, C)$ to $(N, H \\times C)$, then the dimension is reduced to $(N, K)$, where $K$ is the dimension that explains over 90% of the variance and ensures same latent sizes.\n- **Normalization** and **Fitting**: We standardize the latent data and then fit the GMM model. \n\n> Difference of ablation study in Figure 2 and Table 1a, Why not add VAVAE in table 1 comparison too?\n\nFigure 2 and Table 1 are under exact the same settings. \nWe did not include VAVAE mainly because space consideration, i.e. 5 rows for each subtable. \nMAETok do outperform VAVAE, as shown by gFID results in Figure 2 and the SiT-L generation results below:\n\n| Tokenizer  | # Tokens  | LP  | rFID  | gFID  |\n|:---: |:---: |:---: |:---: |:---: |\n| MAETok  | 128  |  72.3 | 0.48  | 5.69  |\n| VAVAE  | 256  |  54.1 |  0.28 |  13.65 |\n\nWe will add this to Table 1 for less confusion.\n\n> Are you able to increase your model size and compare with DC-AE?\n\nWe provide the 512x512 generation results of training 2B USiT with MAETok for 500K steps (as in DC-AE) below:\n\n| Tokenizer  | # Params | # Tokens  | rFID  | gFID w/o CFG  | gFID w/ CFG  |\n|:---:|:---:|:---:|:---:|:---:|:---:|\n| MAETok | 176M | 128  | 0.48  | 1.72 | 1.65 |\n| DC-AE  | 323M | 256  |  0.2s2 |  2.90 | 1.72 |\n\n**MAETok with only 128 tokens established a new SOTA on 512x512 generation**: our gFID w/o CFG already outperforms previous results with CFG.\nWe will include this comparison in our revised paper.\n\n> Ablation analysis on the 2D ROPE?\n\n2D ROPE allows easier and better tokenization with mixed resolution images. \nWe fine-tune MAETok with ROPE and absolute position embedding (APE) on mixed 256x256 and 512x512 images, where 2D ROPE not only presents better results but also generalizes better on higher resolution.\n\n| Position Embedding  | 256x256 rFID  | 512x512 rFID  |\n|:---: |:---: |--- |\n| 2D ROPE  | 0.51  | 0.72  |\n| 2D APE  | 0.73  | 1.43  |\n\n\n> It’s better if you can explain VAVAE a bit before diving into that around Figure 2\n\nWe will add these results in our revised paper.\n\n--- \n\nWe hope the above results could resolve your concern and further validate the effectiveness of MAETok. \nIf you also find the results helpful, please consider raise our rating, thanks!", "claims_and_evidence": "It's a bit contradictive between figure 2 and 4, from Figure 4, AE seems to have fewer GMM mode, i.e. it's more concentrated on one mode.", "methods_and_evaluation_criteria": "What’s the motivation to introduce learnable token z? What would happen if you do not introduce learnable token z?\n\nI think generally this paper proposes several techniques: 1. masked autoencoder, 2. learnable token; 3. auxiliary decoders to align latent feature. I'd like to know what's their weight in the contribution. For example, if you remove 3 and keep 1 and 2, how much quality would drop, and if you remove 1 and keep 2 and 3 what would happen, etc.", "theoretical_claims": "No", "experimental_designs_or_analyses": "I expect the authors to elaborate details on how they get figure 2, what dataset do they use to get statistics of the latent space? Are the latent dimension the same for all the four methods?\n\nWhat's the difference of ablation study in Figure 2 and Table 1a, Why not add VAVAE in table 1 comparison too?\n\nAre you able to increase your model size and compare with DC-AE?\n\nCan you some show ablation analysis on the 2D ROPE, is it useful?", "relation_to_broader_scientific_literature": ".", "essential_references_not_discussed": ".", "other_strengths_and_weaknesses": ".", "other_comments_or_suggestions": "It’s better if you can explain VAVAE a bit before diving into that around Figure 2", "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:

Meta-review: All of reviewers vote for acceptance post-rebuttal. The AC checked all the materials and concurs that the paper has proposed an efficient and effective approach based on masked auto-encoding to train visual tokenizers for diffusion models. The work has made solid contributions to the properties of the latent space that encourages image generation quality, both theoretically and empirically, and therefore should be accepted. Please incorporate necessary changes in the final version, and open-source code/models as promised.
Decision: Accept (spotlight poster)