如果你一直在关注生成式 AI 的爆发，那你一定对潜在扩散模型 (Latent Diffusion Models, LDMs) 并不陌生，它是 Stable Diffusion 等重量级模型背后的架构。LDMs 的秘诀在于效率: 它们不在像素空间逐个像素地生成图像，而是在压缩的“潜在空间 (latent space) ”中进行操作。

这种压缩是由Tokenizer (通常是一个自动编码器) 处理的。多年来，标准的建议一直是使用变分自动编码器 (VAEs) 。 VAE 强制潜在空间服从平滑的高斯分布，理论上这使得扩散模型更容易学习。但这存在一个权衡: 这种平滑性约束通常会导致重建图像模糊，并限制了最终图像的保真度。

如果我们不需要这种约束呢？如果一个标准的自动编码器 (Autoencoder, AE) ——它在保留细节方面表现更好——可以被教导在不使用 VAE 那种繁重的数学约束的情况下组织其潜在空间，结果会怎样？

在最近的一篇论文《Masked Autoencoders Are Effective Tokenizers for Diffusion Models》中，研究人员提出了 MAETok 。 他们证明，通过使用掩码自动编码 (Masked Autoencoding, MAE) 技术，我们可以训练一个普通的 AE 来学习高度结构化、语义丰富的潜在空间。结果如何？仅使用 128 个 Token 就在 ImageNet 上实现了最先进 (SOTA) 的生成效果，同时显著缩短了训练时间并提高了推理吞吐量。

在这篇文章中，我们将解构这篇论文，探讨为什么潜在空间的结构比变分约束更重要，以及掩码机制为何是解锁更强扩散模型的关键。

瓶颈: 扩散模型中的 Tokenizer

在深入解决方案之前，让我们先定义问题。扩散模型的计算成本很高。为了使其具有可扩展性，我们使用一个两阶段的过程:

  1. Token 化 (Tokenization) : 一个自动编码器将图像 \(x\) 压缩成更小的潜在表示 \(z\)。
  2. 生成 (Generation) : 一个扩散模型学习生成这些潜在代码 \(z\)。

最终图像的质量在很大程度上取决于 Tokenizer。如果 Tokenizer 丢失了细节，扩散模型就永远无法恢复它。

VAE 与 AE 的权衡

历史上，研究人员更倾向于使用 VAE 。 VAE 在损失函数中添加了一个正则化项 (KL 散度) ，迫使潜在代码近似标准的正态分布。这确保了潜在空间的平滑性——如果你在已知代码附近采样一个点，它很可能会解码成一个有效的图像。

然而, AE (自动编码器) 通常能实现更好的重建保真度，因为它们不需要对抗 KL 正则化。它们纯粹专注于压缩和解压图像。缺点是什么？它们的潜在空间通常很“混乱”。它们可能高度纠缠或多模态，使得扩散模型极难学习其概率分布。

MAETok背后的研究人员提出了一个关键问题: 对于扩散模型来说，究竟什么样的潜在空间才是“好”的?

理论洞察: 结构胜于正则化

作者假设，VAE 的变分特性并不是帮助扩散模型的关键。相反，关键在于潜在空间的结构 。

为了证明这一点，他们使用高斯混合模型 (GMM) 分析了 AE、VAE 及其提出的 MAETok 的潜在空间。GMM 试图使用高斯曲线的组合来拟合数据分布。

图 2. AE、VAE、VAVAE 和 MAETok 潜在空间上的 GMM 拟合。潜在空间中更少的 GMM 模态通常对应更低的扩散损失。

如图 2 所示，他们发现了一个强相关的联系:

  • 更少的 GMM 模态 (更低的 NLL) \(\rightarrow\) 更低的扩散损失 \(\rightarrow\) 更好的生成质量。

简单来说，如果潜在空间中的数据被分组为清晰、独特的簇 (较少的模态) ，扩散模型就更容易学习生成数据。与标准 AE 相比，VAE 有助于减少这些模态，但它们为此牺牲了图像细节。

直觉背后的数学原理

论文通过理论分析支持了这一实证发现。他们将数据分布定义为 \(K\) 个高斯的混合:

GMM 分布公式

扩散模型 (DDPM) 试图最小化分数匹配损失 (score matching loss) :

分数匹配损失公式

研究人员推导了一个关于样本复杂度的定理——本质上是，为了将模型训练到特定的错误率 (\(\varepsilon\)) ，你需要多少数据 (\(n'\)) 。

样本复杂度公式

仔细看分子。复杂度随 \(K^4\) 缩放。这意味着随着潜在空间中模态数量 (\(K\)) 的增加，训练扩散模型的难度会急剧上升。

结论: 我们不一定需要 VAE。我们只需要一个具有较小 \(K\) (少量模态) 的潜在空间——一个具有判别性 (discriminative) 且分离良好的空间。

解决方案: MAETok

基于“我们需要一个具有判别性的潜在空间，但不需要 VAE 的重建惩罚”这一洞察，作者提出了 MAETok 。

该方法使用了标准的自动编码器架构，但改变了它的训练方式。他们借用了掩码自动编码器 (MAE) 的概念，这一概念在视觉 Transformer 的自监督学习中非常著名。

架构概览

MAETok 采用了基于视觉 Transformer (ViT) 的 1D Tokenizer 设计。

图 3. MAETok 的模型架构。

如图 3 所示，该过程如下运作:

  1. 切分补丁 (Patching) : 输入图像被切分成多个补丁 (patches) 。
  2. 掩码 (Masking) : 这些补丁中的很大一部分 (40-60%) 被随机掩盖掉。
  3. 编码器 (\(\mathcal{E}\)) : 可见的补丁与一组可学习的潜在 Token (latent tokens) 一起被送入编码器。编码器处理这些内容以产生潜在表示 \(h\)。 编码器公式
  4. 解码器 (\(\mathcal{D}\)) : 解码器接收潜在表示和一组掩码 Token 来重建像素值。 解码器公式

为什么要掩码？

为什么掩码对扩散有帮助？当编码器必须仅从 40% 的像素重建图像时，它不能依赖局部的捷径或高频噪声。它被迫去理解图像的全局语义结构 (例如，“这是一只狗，所以这里应该有一条尾巴”) 。

这种语义理解自然导致了一个潜在空间，其中相似的对象被聚集在一起——这正是我们想要的低模态、具有判别性的结构。

辅助解码器: 多任务学习

为了确保编码器学习到丰富的特征，作者不仅仅重建像素。他们引入了辅助浅层解码器 (Auxiliary Shallow Decoders) 。 这些是仅在训练期间使用的小型、临时的解码器，它们试图从被掩盖的 Token 中预测不同的特征，例如:

  • HOG (方向梯度直方图) : 用于边缘和纹理结构。
  • DINOv2 & CLIP: 用于高级语义理解。

辅助损失函数如下所示:

辅助损失公式

其中 \(M\) 是掩码，\(\hat{y}^j\) 是对目标特征 \(j\) 的预测。这种多目标方法确保潜在空间既能捕获视觉细节，又能捕获语义含义。

解耦潜在结构与重建

这里有一个陷阱: 高掩码率对于学习结构非常棒，但它们可能会损害细粒度的像素重建。

为了解决这个问题，作者采用了一种两阶段的训练策略:

  1. 掩码训练: 使用掩码训练编码器和解码器。这构建了结构化的潜在空间。
  2. 解码器微调: 冻结编码器 (保留结构) ，并在未掩码的图像上仅微调像素解码器。

这允许解码器学习如何在编码器提供的稳健结构蓝图上绘制高保真细节。

可视化潜在空间

掩码真的能清理潜在空间吗？让我们看看 UMAP 可视化。

图 4. 潜在空间的 UMAP 可视化。

在图 4 中，对比 AE (a)、VAE (b) 和 MAETok (c)。

  • AE (a) 是一团混乱的重叠颜色——高度纠缠。
  • VAE (b) 更平滑，但仍然有显著重叠。
  • MAETok (c) 显示出清晰、分离的簇。

这证实了掩码迫使编码器在潜在空间中分离不同的概念 (类别) ，从而减少了扩散模型需要学习的模态数量。

此外，空间的“判别性”程度与最终生成质量之间存在直接的相关性。

图 5. 潜在空间质量与生成性能的关系。

图 5(a) 显示，随着线性探测 (LP) 准确率的提高 (这是衡量特征可分性的指标) ，生成 FID (gFID) 也会提高 (越低越好) 。 图 5(b) 显示，与 AE 或 VAE 的潜在空间相比，扩散模型在 MAETok 潜在空间上的训练速度要快得多。

实验结果

研究人员在 ImageNet 的 256x256 和 512x512 分辨率上测试了 MAETok。他们使用了基于 Transformer 的扩散骨干网络 (SiT 和 LightningDiT) 。

定量性能

考虑到效率，结果令人印象深刻。MAETok 仅使用 128 个 Token , 而许多竞争对手使用 256 甚至 1024 个 Token。

表 2. ImageNet 256x256 上的系统级比较。

在表 2 (256x256 分辨率) 中，查看底部的 Ours 行。

  • gFID: MAETok 在不使用无分类器指导 (CFG) 的情况下实现了 2.21-2.31 的 FID，这与使用 VQGAN 或标准 VAE 的模型相当或更好，尽管使用的 Token 更少。
  • 效率: 使用 128 个 Token 极大地降低了扩散 Transformer 的计算成本。

这种优势在 512x512 分辨率下依然存在:

表 3. ImageNet 512x512 上的系统级比较。

表 3 显示，MAETok + SiT-XL 在使用 CFG 的情况下实现了 1.69 的 gFID，击败了参数量大得多的 2B USiT 模型，而自身仅有 675M 参数。

定性结果

数字很好看，但图像效果如何呢？

图 1. MAETok 生成样本。

图 1 展示了在 512x512 分辨率下生成的样本。图像清晰、连贯且纹理丰富。动物的皮毛和水中的倒影证明，128 Token 的压缩并没有丢失关键细节。

以下是更多未经挑选的样本，展示了其多样性:

图 20. 生成的笔记本电脑。

图 22. 生成的跑车。

模型在处理复杂的几何形状 (如笔记本电脑) 和反射表面 (如跑车) 时都保持了高保真度。

讨论与启示

MAETok 的成功挑战了生成式 AI 领域的普遍看法。它表明:

  1. 变分 (Variational) 并非至关重要: 我们不需要 VAE 的 KL 散度来训练好的扩散模型。我们只需要一个易于建模的潜在空间。
  2. 掩码是一种正则化器: 掩码自动编码充当了一种强大的正则化器，自然地构建潜在空间结构，将语义相似的项目分组在一起。
  3. 效率: 我们可以通过高度压缩的潜在表示 (128 个 Token) 实现 SOTA 结果。这对训练速度有巨大的影响。作者报告称，与之前的基准相比, 训练收敛速度快了 76 倍 。

对未来工作的影响

通过将编码器的语义学习 (通过掩码) 与解码器的重建 (通过微调) 解耦，MAETok 提供了一个灵活的框架。未来的工作可以将更强的语义教师 (如大型语言模型或多模态嵌入) 集成到辅助解码器中，以进一步提高生成的一致性。

结论

MAETok 代表了 Tokenizer 设计向前迈出的重要一步。通过摆脱 VAE 并拥抱 MAE 的“掩码与预测”范式，作者创造了一种既简单又更有效的方法。

对于学生和研究人员来说，关键的启示是: 数据表示的结构决定了生成任务的难度。 有时候，组织这种结构的最佳方式不是通过复杂的数学约束，而是通过训练任务本身。

本文基于论文《Masked Autoencoders Are Effective Tokenizers for Diffusion Models》(2025)。