Research map / Image generation / Autoregressive and latent generation
Masked Autoencoders Are Effective Tokenizers for Diffusion Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes MAETok, a plain autoencoder tokenizer trained with masked modeling and auxiliary shallow decoders predicting HOG, DINOv2, and CLIP/SigLIP features, showing that a discriminative (fewer GMM modes) latent space rather than VAE regularization drives better diffusion generation with 128 tokens.
Problem
What latent space structure makes a tokenizer good for diffusion models; existing tokenizers are either VAEs (smooth but lower fidelity) or plain AEs (high fidelity but entangled latents).
Approach
Train a plain 1D transformer autoencoder with mask modeling at the encoder and auxiliary shallow decoders that predict multiple teacher features of masked tokens, yielding a more discriminative, fewer-mode latent space; then freeze the encoder and fine-tune only the pixel decoder to recover reconstruction fidelity. Also supports a theory linking GMM mode count to diffusion sample complexity.
Key results
- MAETok + SiT-XL reaches 1.69 gFID with CFG and 2.79 without CFG at 512x512 with 128 tokens, beating 2B USiT (1.72 w/o CFG in the paper's table row for USiT-2B).
- On 256x256 ImageNet without CFG, SiT-XL on MAETok (128 tokens) achieves 2.31 gFID, outperforming REPA's 5.90 by 3.59.
- MAETok attains rFID 0.48 and SSIM 0.763 at 256x256 with 128 tokens, versus VA-VAE rFID 0.28 with 256 tokens.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use MAETok with 128 tokens instead of the KL tokenizer | SiT-XL on ImageNet 512×512 · ImageNet 512×512 | gFID (w/ CFG): 1.72 -> 1.69 | Helped |
| Add mask modeling to the AE (ablation) | AE tokenizer with SiT-L · ImageNet 256×256 | gFID: 24.47 -> 5.78 | Helped |
| Add mask modeling to the AE (ablation) | AE tokenizer · ImageNet 256×256 | rFID: 0.67 -> 0.85 | Hurt |
| Add decoder fine-tuning after masked AE training | MAETok · ImageNet 256×256 | rFID | Helped |
| Use combined reconstruction targets instead of a single target (ablation) | MAETok with SiT-L · ImageNet 256×256 | gFID: 17.18 -> 5.78 | Helped |
| Use DINO features rather than pixel-only reconstruction (ablation) | MAETok with SiT-L · ImageNet 256×256 | gFID: 17.18 -> 6.24 | Helped |
| Set mask ratio to low=0, high=60 instead of low=40, high=60 (ablation) | MAETok with HOG target and SiT-L · ImageNet 256×256 | gFID: 17.18 -> 24.15 | Hurt |
| Use a 3-block auxiliary decoder rather than a linear decoder (ablation) | MAETok with SiT-L · ImageNet 256×256 | gFID: 6.98 -> 5.78 | Helped |
| Remove mask modeling (ablation) | MAETok · ImageNet 256×256 | gFID: 5.78 -> 8.44 | Hurt |
| Remove learnable tokens (ablation) | MAETok · ImageNet 256×256 | rFID: 0.85 -> 1.01 | Hurt |
| Remove auxiliary decoder (ablation) | MAETok · ImageNet 256×256 | gFID: 5.78 -> 17.18 | Hurt |
| Use CLIP reconstruction target rather than pixel-only target (ablation) | MAETok with SiT-L · ImageNet 256×256 | gFID: 17.18 -> 11.31 | Helped |
Surprising findings
- Variational regularization (KL) is not necessary: plain AE with mask modeling gives better gFID than VAE, which barely improves under masking.
- Freezing the encoder and fine-tuning only the pixel decoder recovers reconstruction fidelity lost to high mask ratios without hurting generation.
Limitations
- CFG tuning is highly sensitive to guidance scale and interval, complicating comparisons
- The theoretical analysis relies on GMM assumptions and loose bounds (K^4 sample complexity) and only informally links to latent structure
- Improvements in CFG-based gFID over strong baselines are smaller than without CFG
Experimental setup
Datasets: ImageNet LAION-COCO
Benchmarks: ImageNet 256x256 ImageNet 512x512 ImageNet reconstruction MS-COCO reconstruction
Metrics: gFID IS rFID PSNR SSIM precision recall linear probing accuracy
Baselines: VA-VAE SD-VAE DC-AE SoftVQ TexTok REPA USiT MAR-H LightningDiT
Builds on: MAE XQ-GAN latent diffusion models SiT LightningDiT REPA
Closest papers by meaning
- Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models (CVPR 2025 · similarity 0.92)
- Boosting Generative Image Modeling via Joint Image-Feature Synthesis (NeurIPS 2025 · similarity 0.92)
- End-to-End Training for Unified Tokenization and Latent Denoising (NeurIPS 2026 · similarity 0.91)
- Stabilize the Latent Space for Image Autoregressive Modeling: A Unified Perspective (NeurIPS 2024 · similarity 0.90)
- Improving the Diffusability of Autoencoders (ICML 2025 · similarity 0.90)
- Nested Diffusion Models Using Hierarchical Latent Priors (CVPR 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.