core

Masked Autoencoders Are Effective Tokenizers for Diffusion Models

Hao Chen, Yujin Han, Fangyi Chen · Carnegie Mellon University, University of Hong Kong, Peking University, AMD · 2025-02 · arXiv:2502.03444

Why it matters

Argues the key variable for generation quality is latent *geometry*, not whether the autoencoder is variational: a masked-modeling-trained, purely discriminative latent with fewer Gaussian-mixture modes is both easier to generate in and reconstructs well, hitting SOTA ImageNet FID with only 128 tokens and far less compute than prior tokenizers.

What this paper does

read: full text

Before this

REPA, VA-VAE and REPA-E had all shown that aligning a latent to a frozen self-supervised encoder helps diffusion training, but each still trained the tokenizer as a variational autoencoder, treating the KL constraint as a given part of the design and leaving the deeper question of what actually makes a latent easy to generate in unexamined.

The problem

The paper asks what property of a latent space actually determines generation quality, and whether the variational, KL-regularized form of the autoencoder is even necessary for that property to hold.

The idea

What matters for diffusion training is the latent distribution's structure, specifically how many Gaussian Mixture modes it takes to fit it and how discriminative its features are, not whether the encoder was trained with a variational constraint. A purely discriminative, non-variational autoencoder trained with masked modeling can produce a latent that is both easier to generate in and reconstructs well.

How it works

MAETok is a ViT-Base encoder-decoder, 176M parameters, mapping images to 128 learnable latent tokens of 32 dimensions each, with no KL loss and no variational bottleneck. During encoding, 40 to 60 percent of image patch tokens are randomly masked and replaced with learnable mask tokens. Alongside ordinary pixel reconstruction (L1, perceptual and adversarial losses), three shallow three-layer auxiliary decoders predict HOG features, frozen DINOv2-Large features and frozen SigCLIP-Large features at the masked positions, so semantic alignment is baked directly into the masked-modeling pretraining rather than added at the diffusion side. After 500K iterations of this mask-modeling stage, the encoder is frozen and only the pixel decoder is fine-tuned for 50K further iterations while the mask ratio decays to zero. The whole tokenizer is trained by ordinary feedforward reconstruction; no autoregressive or diffusion decoding is used anywhere in tokenizer training.

Evidence

MAETok reaches reconstruction FID 0.48 at 256x256 with 128 tokens, beating SoftVQ's 0.61 at twice the token count. Paired with SiT-XL, it reaches gFID 2.31 without guidance and 1.67 with guidance at 256x256, and gFID 1.69 with guidance at 512x512, the paper's headline SOTA number. The authors report this comes with 76x faster tokenizer-plus-generator training and 31x higher inference throughput relative to a REPA-style baseline, with per-image compute dropping from 373.3 to 48.5 GFLOPs. The controlled ablation is direct. With mask modeling and the auxiliary decoders, a plain non-variational autoencoder reaches gFID 5.78, against 24.47 for the same autoencoder without mask modeling, and against a VAE variant that only reaches 18.17 with mask modeling, showing the KL constraint actively works against the effect rather than being neutral to it. Removing the semantic auxiliary targets and keeping only pixel reconstruction leaves gFID at 17.18, showing the DINOv2 and SigCLIP alignment, not masking alone, drives most of the gain.

Limitations

The authors concede that without guidance, their generation numbers still trail some larger autoregressive baselines, and attribute the smaller gain from classifier-free guidance to the latent already being semantically discriminative, which they only partly explore with newer guidance techniques such as autoguidance. High mask ratios of 40 to 60 percent degrade raw reconstruction fidelity on their own, a tradeoff the decoder fine-tuning stage compensates for rather than eliminates. The theoretical account of why fewer Gaussian Mixture modes help, a sample-complexity bound scaling with the fourth power of mode count, is offered as a proxy explanation rather than a full theory of generation quality.

Why it matters

It reframes the whole representation-alignment line. REPA, VA-VAE and REPA-E each read as ways of pushing a variational latent toward a semantic, low-mode structure using an external encoder's signal; MAETok shows the variational form was never the necessary ingredient; the structure is, and it can be trained directly with masked modeling and feedforward reconstruction, with no diffusion or autoregressive decoding required in the tokenizer at all.

Abstract, in the authors' own words

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.

Research line

Representation-aligned latentsascendant

Reported results

Benchmark	Value	Guidance	Budget	Source
ImageNet 256x256 gFID	2.31	none	4M steps	Table 2
ImageNet 256x256 gFID	1.67	cfg=2.0, 250 steps	4M steps	Table 2
ImageNet 256x256 rFID	0.48	—	—	Table 4

Design-axis choices

Representation

Semantic / foundation latent

Problem

Semantic / foundation-model latents discard much of the high-frequency pixel detail (exact color, texture, fine structure) that faithful reconstruction — and, later, edit-region preservation — depends on.

Method note — the shared flow-matching interpolation

Every flow-matching / rectified-flow paper in this atlas trains toward a straight-line path between a noise sample x₀ and a data sample x₁:

x_t = (1-t)\,x_0 + t\,x_1, \qquad v_\theta(x_t, t) \approx x_1 - x_0

Builds on

Nothing recorded yet.

Built on by

Nothing recorded yet.

Challenges / competes with

challenges VAE latent — MAETok shows the variational form of the autoencoder is not necessary; a purely discriminative, masked-modeling-trained latent with fewer Gaussian-mixture modes generates better.
competes_with Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models — DC-AE and MAETok attack the same latent-design problem from different angles: architecture and compression ratio (DC-AE) versus training objective and latent geometry (MAETok).