仅128个token达到ImageNet生成SOTA性能！MAETok:有效的扩散模型的关键是什么？(卡内基梅隆&港大等)

论文链接：https://arxiv.org/pdf/2502.03444
项目链接：https://github.com/Hhhhhhao/continuous_tokenizer

亮点直击

• 理论与实验分析：通过实验和理论分析建立了隐空间结构与扩散模型性能之间的联系。揭示了具有更少高斯混合模型（GMM）模式的结构化潜空间能够更有效地训练和生成扩散模型。
• MAETok：使用mask建模训练普通自编码器（AE），并证明了具有更具判别性的潜空间的简单 AE 能够加速学习、提高生成质量，并提升扩散模型的吞吐量，从而表明 VAE 的变分正则化并非必要。
• SOTA 生成性能：在 MAETok 训练的扩散模型（675M 参数，128 token）在 256×256 ImageNet 生成任务上达到了与先前最佳模型相当的性能，并在 512 分辨率上超越了 2B USiT，取得了 1.69 gFID 和 304.2 IS。

结果展示

总结速览

解决的问题

现有的扩散模型通常使用变分自编码器（VAE）作为tokenizer，但VAE的变分约束可能限制了重建质量，影响生成效果。此外，现有的自动编码器（AE）虽然能提供更高保真度的重建，但其隐空间可能过于混乱，不利于扩散模型的训练和采样。因此，如何构建一个既能保持高重建质量，又能提供良好结构化隐空间的tokenizer仍然是一个未解决的问题。

提出的方案

本文提出MAETok，一种基于 mask 建模（Masked Autoencoder, MAE）的自编码器（AE）。它通过随机屏蔽图像token并重建缺失部分，以学习更加语义丰富、结构更清晰的隐空间。核心设计包括：

• mask建模：在编码器端随机屏蔽一部分图像token，要求解码器恢复原始特征，以迫使AE学习更具判别性的特征表示。
• 辅助浅层解码器：在标准像素解码器之外，增加辅助解码器，用于预测latent token的特征（如HOG、DINOv2、CLIP等），增强隐空间表示的语义信息。
• 编码器冻结策略：在高 mask 比（40%-60%）训练后，固定编码器并微调解码器，以恢复像素级重建质量，同时保留高质量的隐空间表示。

应用的技术

• 扩散模型（Diffusion Models）
• Mask自编码器（MAE）进行特征学习
• 高效Transformer架构用于tokenizer设计
• 高 mask 比训练策略
• 隐空间分析（使用高斯混合模型GMM）优化扩散模型的学习

达到的效果

• 更优的隐空间：隐空间特征更加判别清晰，且具有更少的GMM模式，提升扩散模型的生成质量。
• 高效训练与推理：相比基于VAE的扩散模型，训练速度提升76倍，推理吞吐量提升31倍（512×512图像）。
• 优越的生成质量：在ImageNet 256×256 和 512×512 基准上，MAETok 使用仅128个token即可实现最先进（SOTA）的生成质量，gFID 达到 1.69。
• 突破VAE限制：证明高质量的隐空间比变分约束（如VAE）更关键，实现无需VAE的扩散模型训练。

方法

VAE 的变分形式对于扩散模型可能并非必要，简单的 AE 便足以在 128 token的条件下实现 SOTA 生成性能，只要它们具有判别性的潜空间，即具有更少的 GMM 模式。本文将该方法称为 MAETok，其详细信息如下。

结构

在最近的 1D 分词器 设计基础上构建 MAETok，该分词器采用可学习的latent tokens。编码器 (E) 和解码器 (D) 都采用 Vision Transformer (ViT) 结构，但经过调整以同时处理 图像tokens 和 latent tokens，如下图 3 所示。

Mask 建模

像素解码器微调

虽然 Mask 建模促使编码器学习更好的隐空间，但较高的 Mask 比例可能会降低即时重建效果。为了解决这一问题，在使用 Mask 建模训练自编码器之后，冻结编码器，从而保留隐空间表示，并仅对像素解码器进行少量额外训练。这一过程使解码器能够更紧密地适应冻结的干净图像隐空间编码，恢复在 Mask 训练过程中丢失的细节。在像素解码器微调过程中使用与公式 (6) 相同的损失，并在此阶段丢弃所有辅助解码器。

实验

本文进行全面的实验来验证 MAETok 的设计选择，分析其隐空间，并基准测试生成性能，展示其优越性。

实验设置

扩散模型的实现细节
在训练 MAETok 之后，使用 SiT和 LightningDiT来执行基于扩散的图像生成任务。将 Patch 大小设置为 1，并使用 1D 位置 embedding，其他参数遵循它们的原始训练设置。使用 458M 参数的 SiT-L 进行分析和消融研究。对于主要结果，训练了 675M 参数的 SiT-XL，经过 4M 步，和 LightningDiT 训练 400K 步，在分辨率为 256 和 512 的 ImageNet 上进行训练。

评估  

对于 Tokenizer 评估，报告了 ImageNet 和 MS-COCO验证集上的重建 Fréchet Inception Distance（rFID）、峰值信噪比（PSNR）和结构相似性指数（SSIM）。对于 Tokenizer 的潜在空间评估，我们对平坦的潜在表示进行线性探测（LP），并报告其准确性。为了评估生成任务的性能，报告了生成 FID（gFID）、Inception Score（IS）（Salimans et al., 2016）、精度和召回率（Precision and Recall），有无无分类器引导（CFG），使用 250 次推理步骤。

MAETok的设计选择

首先进行了一项广泛的消融研究，以了解Mask建模和不同设计如何影响tokenizer的重建效果，更重要的是，如何影响扩散模型的生成效果。从一个AE开始，逐步添加不同的组件，研究AE的rFID和SiT-L的gFID。

Mask建模
在下表1a中，我们比较了AE和VAE的Mask建模效果，并研究了像素解码器微调的提议。对于AE，Mask建模显著提高了gFID，并略微降低了rFID，这可以通过解码器微调阶段恢复，而不会牺牲生成性能。相反，Mask建模对VAE的gFID仅有轻微改善，因为施加的KL约束可能会妨碍潜在空间的学习。

重建目标
在上表1b中，我们研究了不同的重建目标如何影响Mask建模中的潜在空间学习。我们展示了使用低级重建特征，如原始像素（仅使用像素解码器）和HOG特征，已经能够学习到更好的潜在空间，从而实现更低的gFID。采用语义教师，如DINO-v2和CLIP，则可以显著提高gFID。结合不同的重建目标可以在重建保真度和生成质量之间实现平衡。

Mask比例
在上表1c中，我们展示了合适的Mask比例对于使用HOG目标学习潜在空间的重要性，正如前人工作中所强调的。较低的Mask比例会阻止AE学习到更具区分性的潜在空间。较高的Mask比例则在重建保真度和潜在空间质量之间存在权衡，从而影响生成性能。

辅助解码器深度
在上表1d中，我们研究了具有多重重建目标的辅助解码器的深度。我们发现，解码器过浅或过深都会损害重建保真度和生成质量。当解码器过浅时，组合目标特征可能会混淆潜在空间中的高级语义和低级细节，从而导致更差的重建保真度。然而，过深的辅助解码器可能会凭借其强大的能力学习到AE的潜在空间中较少的区分特征，从而也导致更差的生成性能。

潜在空间分析

本文进一步分析了AE变体的潜在空间与SiT-L生成性能之间的关系。

潜在空间可视化
提供了UMAP可视化如下图4所示，以直观地比较不同AE变体所学习到的潜在空间。值得注意的是，AE和VAE展示了更为纠缠的潜在 embedding，其中不同类别的样本往往有较大的重叠。MAETok展示了明显分离的聚类，并且类别之间的边界较为清晰，这表明MAETok学习到了更具区分性的潜在表示。与下图2中的分析一致，MAETok的潜在表示更加区分且分离，导致了更少的GMM模式，并提升了生成性能。

潜在分布与生成性能
我们通过研究潜在空间上的线性探测（LP）准确度，作为潜在代码中语义信息保留程度的代理，并与生成性能的gFID进行关联，来评估潜在空间的质量。在图5a中，我们观察到，潜在分布更具区分性的tokenizer，表现在更高的LP准确度，相应地也能实现更低的gFID。这个发现表明，当特征在潜在空间中良好聚类时，生成器能够更容易学习生成高保真样本。我们进一步通过跟踪训练过程中gFID的变化来验证这一直觉，如下图5b所示，MAETok使得训练更快收敛，gFID迅速下降，且低于AE或VAE基准。高质量的潜在分布被证明是实现强大的最终生成指标并加速训练的关键因素。

主要结果

生成  

在下表2和表3中，分别基于256×256和512×512的ImageNet基准，比较了SiT-XL和LightningDiT的MAETok变体，并与其他SOTA生成模型进行了比较。值得注意的是，使用仅128个token和普通AE架构训练的SiT-XL，在没有使用CFG的情况下，始终能取得更好的gFID和IS：它在256分辨率下比REPA提高了3.59的gFID，并在512分辨率下达到了2.79的SOTA可比gFID。当使用CFG时，SiT-XL与基于VAEs的竞争性自回归和扩散基准在256分辨率下的表现相当。它击败了使用256个token的2B USiT，并且在512分辨率下实现了新的SOTA，gFID为1.69，IS为304.2。使用更多先进技巧训练的LightningDiT表现更好，它在没有CFG的情况下超越了1B参数的MAR-H和2B参数的USiT，gFID为2.56，IS为224.5，并且在使用CFG时，gFID为1.72。这些结果表明，潜在空间的结构（见上图4），而非tokenizer的变分形式，对于扩散模型的有效学习至关重要。

重建  

MAETok在ImageNet和MS-COCO上的重建能力也非常强，如下表4所示。与之前的连续tokenizer（包括SD-VAE、DC-AE、VA-VAE、SoftVQ-VAE和TexTok）相比，MAETok在重建质量和潜在空间大小之间实现了一个良好的平衡。在256×256的ImageNet上，使用128个token，MAETok的rFID为0.48，SSIM为0.763，超越了SoftVQ等方法，在保真度和感知相似度方面均表现优越，同时使用了TexTok一半的token数。在MS-COCO上，尽管tokenizer没有直接训练，MAETok仍能提供强大的重建能力。在512分辨率下，MAETok通过平衡压缩比和重建质量，保持了其优势。

讨论

高效的训练与生成  

1D tokenizer设计的一个显著优势是，它能够支持任意数量的潜在token。通常256×256和512×512图像被编码为256和1024个token，而MAETok在两种情况下都使用128个token。这大大提高了扩散模型的训练和推理效率。例如，在使用512×512图像的1024个token时，SiT-XL的Gflops和推理吞吐量分别为373.3和每秒0.1张图像。而MAETok将Gflops降至48.5，将吞吐量提高到每秒3.12张图像。通过改进的收敛性，MAETok使得训练速度比REPA快了76倍，表现相当。

无条件生成  

从我们的结果中得到的一个有趣观察是，与以前的方法相比，使用MAETok训练的扩散模型在没有CFG的情况下通常展现出显著更好的生成性能，同时在使用CFG时性能差距较小。我们推测原因是，无条件类别也学习了潜在空间中的语义，如表5中的无条件生成性能所示。随着潜在空间变得更加区分性，无条件生成性能也显著提高。这意味着CFG线性组合方案可能变得不那么有效，这与我们在附录C.2中提供的CFG调优结果一致。

结论

本文对扩散模型的隐空间特性进行了理论和实证分析，证明了隐空间分布中较少的模式有助于更有效的学习和更好的生成质量。基于这些见解，开发了MAETok，它通过mask建模实现了最先进的性能，而无需使用变分约束。仅使用128个token，本文的方法显著提高了ImageNet上的计算效率和生成质量。研究结果表明，具有更强区分性的隐空间，而非变分约束，是有效的扩散模型的关键，为大规模高效生成建模开辟了新的方向。

本文转自AI生成未来 ，作者：AI生成未来

原文链接:​​https://mp.weixin.qq.com/s/ZWhWR_Sc275SDl1BRVCH2g​​