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
分类: cs.CV, cs.AI, cs.LG
发布日期: 2025-02-05 (更新: 2025-05-30)
💡 一句话要点
提出MAETok，利用掩码自编码器为扩散模型学习更优的token表示，显著提升图像生成质量和效率。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 掩码自编码器 扩散模型 图像生成 潜在空间 tokenizer 生成模型 自编码器
📋 核心要点
- 现有潜在扩散模型tokenizer的潜在空间性质研究不足，限制了生成质量的进一步提升。
- 提出MAETok，利用掩码自编码器学习语义丰富的潜在空间，无需变分约束，提升生成质量。
- 实验表明，MAETok在ImageNet生成上达到SOTA，gFID为1.69，训练速度提升76倍，推理吞吐量提升31倍。
📝 摘要（中文）
最近，潜在扩散模型在高分辨率图像合成方面表现出卓越的性能。然而，对于扩散模型而言，tokenizer的潜在空间的性质对于更好的学习和生成仍然缺乏深入研究。本文从理论和实验两方面发现，生成质量的提升与具有更好结构的潜在分布密切相关，例如具有更少高斯混合模式和更具区分性的特征的潜在分布。受此启发，我们提出了MAETok，一种利用掩码建模的自编码器（AE），以学习语义丰富的潜在空间，同时保持重建保真度。大量的实验验证了我们的分析，表明自编码器的变分形式不是必需的，并且仅来自AE的判别潜在空间就能够在使用仅128个token的情况下实现ImageNet生成的最新性能。MAETok实现了显著的实际改进，在512x512生成中实现了1.69的gFID，训练速度提高了76倍，推理吞吐量提高了31倍。我们的研究结果表明，潜在空间的结构，而不是变分约束，对于有效的扩散模型至关重要。代码和训练好的模型已发布。
🔬 方法详解
问题定义：现有潜在扩散模型依赖于tokenizer将图像压缩到潜在空间，但tokenizer的潜在空间结构对生成质量的影响缺乏深入研究。现有方法，如VQ-VAE，通常采用变分自编码器，但变分约束可能会限制潜在空间的表达能力，导致生成质量受限。因此，如何设计一个能够学习到更具判别性和语义信息的潜在空间的tokenizer，是本文要解决的关键问题。
核心思路：本文的核心思路是利用掩码自编码器（MAE）学习潜在空间表示。MAE通过随机mask输入图像的部分区域，并重建被mask的部分，从而迫使模型学习到图像的整体结构和语义信息。与传统的变分自编码器不同，MAE不施加变分约束，从而允许模型学习到更具表达能力的潜在空间。这种更具判别性的潜在空间能够更好地支持扩散模型的学习和生成。
技术框架：MAETok的整体框架包括一个编码器和一个解码器。编码器将输入图像编码为潜在表示，解码器则从潜在表示重建图像。关键在于编码器部分采用了掩码建模策略，即随机mask输入图像的部分区域，并将未被mask的部分输入到编码器中。解码器则需要根据编码器的输出重建完整的图像，包括被mask的部分。扩散模型则在MAETok学习到的潜在空间上进行训练和生成。
关键创新：最重要的技术创新点在于使用掩码自编码器作为扩散模型的tokenizer。与传统的变分自编码器相比，MAE能够学习到更具判别性和语义信息的潜在空间，从而提升扩散模型的生成质量。此外，本文还发现，变分约束对于扩散模型的tokenizer来说不是必需的，甚至可能是有害的。
关键设计：MAETok的关键设计包括掩码比例、掩码策略、损失函数和网络结构。掩码比例控制了输入图像被mask的区域的比例，通常设置为较高的值（如75%）以迫使模型学习到更多的上下文信息。掩码策略可以是随机的，也可以是基于图像内容的。损失函数通常采用重建损失，即衡量重建图像与原始图像之间的差异。网络结构可以采用Transformer或卷积神经网络等。
🖼️ 关键图片
📊 实验亮点
MAETok在ImageNet 512x512图像生成任务上取得了显著的性能提升，gFID达到1.69。与现有方法相比，MAETok的训练速度提高了76倍，推理吞吐量提高了31倍。实验结果表明，MAETok能够学习到更优的潜在空间表示，从而提升扩散模型的生成质量和效率。
🎯 应用场景
MAETok具有广泛的应用前景，可用于图像生成、图像编辑、图像修复等领域。通过学习更优的潜在空间表示，MAETok可以提升生成图像的质量和效率，降低计算成本。此外，MAETok还可以应用于其他生成模型，如GAN等，提升其性能。
📄 摘要（原文）
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.