当前位置：首页 > news >正文

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

news 来源：原创 2025/5/9 22:04:10

论文链接：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 所示。

编码器。编码器首先根据预定义的块大小P将输入图像划分为个块，每个块映射到一个维度为的 embedding 向量，从而得到image tokens 。此外，定义一组个可学习的latent tokens 。编码器的变换器将image patch embeddings和latent tokens的连接作为输入，并仅从latent tokens输出潜在表示，其中为其维度：

解码器。为了重建图像，使用一组个可学习的image tokens 。将这些mask tokens与连接，作为解码器的输入，并仅使用来自mask tokens的输出进行重建。

在上使用一个线性层来回归像素值，并获得重建的图像。

位置编码。为了编码空间信息，在编码器的mage patch tokens和解码器的image tokens 上应用二维旋转位置 embedding （RoPE）。相比之下，latent tokens （及其编码对应）使用标准的1D绝对位置 embedding ，因为它们不映射到特定的空间位置。这种设计确保了patch-based tokens保留了二维布局的概念，而学习到的latent tokens 被视为transformer架构中的一组抽象特征。

训练目标。使用与先前工作（Esser et al., 2021）相同的标准tokenizer损失函数来训练MAETok。

其中、和分别表示像素级均方误差（MSE）损失、感知损失和对抗损失，和是超参数。请注意，MAETok是一个简单的自编码器架构，因此，它不像变分自编码器（VAE）那样需要后验和先验之间的变分损失，这简化了训练过程。

Mask 建模

编码器的Token Mask。 MAETok的一个关键特性是，在训练过程中引入了 mask 建模，遵循MAE的原则，以自监督的方式学习更具辨别性的潜在空间。根据二进制masking指示器，随机选择一定比例（例如40%-60%）的image patch tokens，并用可学习的mask token替换它们，然后将其输入编码器。所有潜在token保持不变，以便更强地聚合未unmasked image tokens的信息，并在解码器输出时用于重建masked tokens。

辅助浅层解码器。在MAE中，需要一个浅层解码器或一个线性层来从剩余的图像token中预测目标特征，例如原始像素值、HOG特征和来自预训练模型的特征。然而，由于我们的目标是将MAE训练为token化器，因此像素解码器需要能够高保真地重建图像。因此，将的容量保持与相似，并结合辅助浅层解码器来预测额外的特征目标，这些解码器与主像素解码器设计相同，但层数较少。形式上，每个辅助解码器接受潜在表示并与它们自己的拼接作为输入，输出作为其特征目标的重建。

其中表示目标特征的维度。与AE一起训练这些辅助解码器，并仅根据 mask 指示器对masked tokens应用额外的MSE损失，类似于Xie et al.。

像素解码器微调

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

实验

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

实验设置

Tokenizer的实现细节。使用XQ-GAN代码库来训练MAETok。使用ViT-Base，从头初始化，用于编码器和像素解码器，它们总共有176M个参数。设置潜在空间为和。三个MAETok变体分别在256×256的ImageNet、512×512的ImageNet和512×512的LAION-COCO子集上训练500K迭代。在第一阶段，使用ImageNet进行mask建模训练时，采用了一个40-60%的mask比例，这是通过消融实验设置的，并且有三个辅助浅层解码器，分别用于多个目标特征：HOG、DINO-v2-Large和SigCLIP-Large。对于在LAION-COCO上训练的tokenizer，采用额外的辅助解码器，用于预测图像的文本描述的离散索引，使用BPE tokenizer。每个辅助解码器有3层，也是通过消融实验设置的。设置和。对于像素解码器微调，在线性地减少 mask 比例，从60%减少到0%，并进行50K迭代，使用相同的训练损失。

扩散模型的实现细节
在训练 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上的计算效率和生成质量。研究结果表明，具有更强区分性的潜在空间，而非变分约束，是有效的扩散模型的关键，为大规模高效生成建模开辟了新的方向。

参考文献

[1] Masked Autoencoders Are Effective Tokenizers for Diffusion Models