PaperlayerView exact revision on arXiv

Source-linked AI summary

Masked Autoencoders Are Effective Tokenizers for Diffusion Models

Hao Chen, Yujin Han, Fangyi Chen, Xiang Li, Yidong Wang, Jindong Wang, Ze Wang, Zicheng Liu, Difan Zou, Bhiksha Raj

arXiv:2502.03444v2cs.CVcs.AIcs.LG

TL;DR

The paper asks what makes a tokenizer latent space effective for diffusion models, a question complicated by the trade-off between VAE reconstruction fidelity and AE latent organization. It analyzes latent-space structure and proposes MAETok, whose mask-modeled plain autoencoder achieves state-of-the-art ImageNet generation with 128 tokens, including 1.69 gFID and substantially faster training and inference. The authors conclude that discriminative latent structure matters more than variational constraints.

  • Problem

    The properties of tokenizer latent spaces that support effective diffusion learning and generation remain under-explored, amid a trade-off between VAE fidelity and AE latent organization.

  • Method

    MAETok trains a plain autoencoder with encoder mask modeling and auxiliary shallow decoders that predict multiple masked-token features while preserving pixel reconstruction.

  • Results

    MAETok achieves state-of-the-art ImageNet generation with 128 tokens, including 1.69 gFID at 512×512, 76× faster training, and 31× higher inference throughput.

  • Takeaways & Limitations

    Latent-space structure, particularly discriminative organization and fewer modes, is more crucial for effective diffusion models than variational constraints.

  • Takeaways & Limitations

    The analysis relies on assumptions concerning data separability, warm-start initialization, and bounded target-distribution means, while the broader impact of synthetic media remains a concern.

Abstract

from arXiv · show

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 76x faster training and 31x higher inference throughput for 512x512 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 are released.

1. Introduction

The paper argues that diffusion quality depends on latent-space structure, especially discriminative organization and fewer GMM modes, rather than variational constraints. It introduces MAETok, which uses mask modeling with plain autoencoders to retain reconstruction fidelity while improving diffusion performance and efficiency.

  • MAETok: MAETok trains plain autoencoders with encoder mask modeling and auxiliary decoders that predict masked HOG, DINO-v2, and CLIP features.The method uses 40–60% masking while retaining a pixel decoder for reconstruction.
  • Latent-space motivation: Latent spaces with fewer Gaussian Mixture Model modes and more discriminative features support more effective diffusion training and generation.The paper connects latent-space structure to diffusion losses and generation quality through theoretical and empirical analysis.
  • Results: Diffusion models trained on MAETok with 128 tokens achieve 1.69 gFID and 304.2 IS at 512 resolution, outperforming the previous 2B-parameter USiT model.The reported diffusion model has 675M parameters.

2. On the Latent Space and Diffusion Models

The paper studies how tokenizer latent distributions affect diffusion learning by comparing their GMM structure and diffusion behavior. Both experiments and theory indicate that fewer latent-space modes reduce learning difficulty, training requirements, and improve generation quality.

  • Empirical Analysis: The study compares AE, VAE, and VAVAE tokenizers by fitting Gaussian mixture models to their latent spaces and training diffusion models under matched conditions.The comparison links latent-space properties to final image-generation quality.
  • Empirical Analysis: Fewer GMM modes correspond to lower diffusion losses and better gFID, with more separated and discriminative latent features reducing learning difficulty.VAVAE requires fewer modes than AE and VAE for comparable GMM fitting quality.
  • Theoretical Analysis: The theoretical analysis models latent data as a Gaussian mixture and shows that more modes typically require larger diffusion-model training sample sizes.The result targets comparable generation quality under the stated assumptions.
  • Theoretical Analysis: For comparable O(Tϵ2) generation quality, the required training sample size scales as O(K4) with the number of mixture modes K.This helps explain why AE and VAE latent spaces can yield worse generations and higher gFID with finite data.
  • MAETok Architecture: MAETok uses a plain 1D ViT autoencoder tokenizer with mask modeling and auxiliary shallow decoders to predict features of masked tokens.The architecture uses a 40–60% mask ratio and multiple target features.

3. Method

MAETok is a plain autoencoder tokenizer that uses masked image-token modeling to learn discriminative latent representations while preserving reconstruction fidelity. Its design combines transformer-based image and latent tokens, auxiliary feature decoders, standard tokenizer losses, and decoder-only fine-tuning.

  • Masked Autoencoder Tokenizer: MAETok trains a plain autoencoder with mask modeling, avoiding variational loss while targeting a more discriminative latent space.The method masks image patch tokens before encoding and uses the resulting representations for tokenizer learning.
  • Architecture: The encoder and decoder use Vision Transformers adapted to process image tokens and learnable latent tokens.The encoder concatenates image patch embeddings with learnable latent tokens and outputs latent representations from the latent-token positions.
  • Training Objectives: The training objective combines pixel-wise reconstruction, perceptual, and adversarial losses weighted by hyperparameters.MAETok uses standard tokenizer losses without a variational posterior-prior regularization term.
  • Auxiliary Shallow Decoders: Auxiliary shallow decoders predict additional feature targets from latent representations and decoder tokens, with losses applied only to masked tokens.These auxiliary decoders retain similar design to the main pixel decoder but use fewer layers.
  • Decoder Fine-Tuning: Freezing the encoder and fine-tuning only the pixel decoder recovers reconstruction details while preserving latent representations learned under masking.This adaptation is performed for a small number of additional epochs on clean images.

4. Experiments

Experiments show that mask modeling gives plain autoencoders discriminative latent spaces while preserving reconstruction quality, enabling strong and efficient ImageNet generation with 128 tokens.

  • Design Choices: Mask modeling significantly improves AE generation quality, while decoder fine-tuning recovers the associated reconstruction-fidelity loss.For VAE, mask modeling yields only marginal gFID improvement, likely because the KL constraint hinders latent-space learning.
  • Design Choices: Semantic reconstruction targets such as DINO-v2 and CLIP improve gFID more than low-level targets, while combining targets balances fidelity and generation quality.Raw pixels and HOG features already improve latent-space quality compared with the baseline.
  • Latent Space Analysis: MAETok produces separated class clusters, unlike the more entangled AE and VAE embeddings, indicating a more discriminative latent representation with fewer GMM modes.The separated representation is associated with improved generation performance.
  • Latent Space Analysis: Higher latent-space linear-probing accuracy corresponds to lower gFID, suggesting that well-clustered features make diffusion training easier.The reported relationship links semantic information preserved in latent codes with generation performance.
  • Discussion: MAETok reduces SiT-XL computation from 373.3 to 48.5 Gflops and raises throughput from 0.1 to 3.12 images/second for 512×512 generation.The paper reports 76× faster training to perform similarly to REPA.
  • Design Choices: MAETok training combines learnable tokens, mask modeling, and auxiliary decoders, with complementary gains across reconstruction and generation quality.Removing masking or the auxiliary decoder weakens results.

5. Related Work

Related work develops image tokenizers and generative models around compact latent representations, reconstruction fidelity, semantic understanding, and scalable generation.

  • Image Tokenization: Image tokenization transforms high-dimensional images into compact, structured latent representations, beginning with autoencoders and extending to variational and vector-quantized methods.Recent work focuses on improving compression fidelity and scalability.
  • Image Tokenization: Recent tokenizers seek to combine high-fidelity reconstruction with semantic understanding within a single representation.This direction builds on advances in compression fidelity and scalability.
  • Image Generation: Image-generation research mainly comprises autoregressive and diffusion paradigms, with Transformer architectures improving autoregressive scalability and diffusion models showing strong performance.The passage describes both paradigms as major directions in image generation.

6. Conclusion

The paper concludes that latent-space structure, rather than variational constraints, is central to effective diffusion learning and generation.

  • Conclusion: Fewer modes in latent distributions enable more effective diffusion learning and better generation quality.MAETok applies this insight through mask modeling without requiring variational constraints.
  • Conclusion: Using 128 tokens, MAETok improves computational efficiency and ImageNet generation quality while achieving state-of-the-art performance.The conclusion presents discriminative latent space as the key factor behind effective diffusion models.

Impact Statement

The work improves the efficiency and effectiveness of image-generation systems while acknowledging broader societal implications of image-synthesis advances.

  • Impact Statement: The paper focuses on more efficient and robust image generation, with potential benefits for creative tools and design.These broader applications are presented alongside concerns about synthetic media.
  • Impact Statement: Synthetic-media advances may raise broader societal concerns, motivating ongoing discussion about responsible deployment.The statement does not specify particular deployment policies or safeguards.

A. Theoretical Analysis

The theoretical analysis connects latent-space structure to diffusion training and sampling, focusing on Gaussian-mixture mode count under boundedness and separation assumptions. It concludes that more modes increase the data and training requirements for comparable estimation or sampling quality.

  • Latent-space model: The analysis models tokenizer latents with Gaussian mixtures and studies how mode count K affects diffusion learning.The framework assumes equally weighted Gaussian components and examines score matching and DDPM behavior.
  • Assumptions: The theory assumes separated mixture means, warm initialization, and bounded maximum mean norm B.These assumptions constrain data geometry, optimization initialization, and target-distribution regularity.
  • Interpretation: Because different latent spaces have similar mean-norm upper bounds, the analysis primarily attributes differences in difficulty to the number of modes K.Figure 8 supports treating B as comparable across AE, VAE, VAVAE, and MAETok.
  • Training error: More mixture modes require more training samples to achieve the same estimation error.The result follows from the DDPM estimation theorem under the stated assumptions.
  • Sampling procedure: Early stopping is analyzed over t ∈ [0, 0.8], and exponentially decreasing step sizes can make the discretization quantity Π = O(1/N^2).This provides a sampling-error control condition within the theoretical setup.
  • Sampling error: The analysis connects training error to sampling error through a KL-divergence bound for the generated distribution.The bound includes terms for diffusion time, score-estimation error, discretization, and the mode-dependent sample requirement.

B.3. Training Details of GMM Models

The GMM experiments standardize latent representations before fitting models to compare tokenizer latent distributions. They use PCA, normalization, and full-ImageNet fitting under a documented computational setup.

  • GMM preprocessing: Latents are flattened, reduced with PCA to exceed 90% explained variance, standardized, and then fit with GMMs.The preprocessing gives all tokenizer latent spaces consistent dimensions before likelihood evaluation.
  • Tokenizer comparison: The comparison trains AE, KL-VAE, and MAETok under identical settings while using a pretrained VAVAE.All tokenizers are evaluated with the same latent size and input dimensions.
  • Computational setup: GMM fitting uses the full ImageNet dataset, batch size 256, and a single NVIDIA A8000.The reported fitting times for 50, 100, and 200 components are roughly 3, 8, and 11 hours.
  • Diffusion training: The SiT-L models used for latent-space generation comparisons train for 400K iterations with AdamW, learning rate 1e-4, and no weight decay.These settings are used for each tokenizer latent space in the comparison.

C.1. More Quantitative Generation Results

Additional quantitative evaluations report precision and recall for generation on the 256×256 and 512×512 ImageNet benchmarks.

  • ImageNet evaluation: Precision and Recall are evaluated on the 256×256 and 512×512 ImageNet benchmarks.The results are provided in Table 10 and Table 11, respectively.

C.2. Classifier-free Guidance Tuning Results

Classifier-free guidance requires careful tuning because small guidance-scale changes can substantially affect gFID. The authors use interval guidance and report that MAETok’s unconditional class learns semantics, complicating the linear CFG scheme.

  • CFG sensitivity: gFID changes significantly even with small classifier-free guidance scales, making guidance tuning difficult.Applying a CFG interval can mitigate the issue by removing guidance at high timesteps.
  • Reported settings: The reported settings use guidance scales from 1.5 to 1.9 with timestep intervals tailored to resolution and diffusion model.For 256×256 models, both intervals are [0, 0.75]; 512×512 models use intervals [0, 0.7] or [0, 0.65].
  • Interpretation: The authors attribute CFG-tuning difficulty to semantics learned by the unconditional class.They leave more advanced CFG schemes as future work.
  • Latent-space analysis: MAETok variants generally learn more discriminative latent spaces with fewer GMM models across reconstruction targets.The comparison is shown through additional UMAP visualizations.
  • Token ablation: Using 128 latent tokens achieves similar reconstruction and downstream generation performance to using 256 tokens.Learnable latent tokens outperform using image tokens only, while 2D RoPE improves mixed-resolution generalization.

C.5. More Qualitative Generation Results

This section presents uncurated MAETok + SiT-XL generation results across ImageNet classes at 256×256 and 512×512 resolutions. It also provides system comparisons, CFG tuning, latent-token and 2D RoPE ablations, and training configuration details.

  • Ablations and configuration: Training configurations for LightningDiT on 256×256 and 512×512 ImageNet are provided.The configuration is documented separately from the qualitative and ablation results.
  • Qualitative generation results: Uncurated 256×256 generations are shown for Loggerhead, Macaw, Cacatua galerita, and Flamingo using CFG 3.0.The corresponding figures identify each ImageNet class label and class index.
  • Quantitative comparisons: System-level comparisons report tokenizer and generator models, parameter counts, latent-token counts, rFID, gFID, Precision, and Recall for 256×256 and 512×512 ImageNet generation.The tables distinguish classifier-free-guidance and no-guidance settings and define the model and parameter-count columns.
  • Ablations and configuration: Additional experiments tune CFG for SiT-XL and ablate latent-token counts and 2D RoPE with MAETok on 256×256 ImageNet.The CFG study uses 10K generated samples, while the ablation reports tokenizer rFID and unguided SiT-L gFID under specified training schedules.
  • Qualitative generation results: Uncurated 512×512 generations are shown for eight classes, including Siberian husky, Arctic fox, Otter, Guitar, and Fire Truck, using CFG 2.0.Additional examples include Go-kart, Laptop, Carriage, and Sports Car.
Loading 2502.03444v2…