MicroDiffusion——采用新的掩码方法和改进的 Transformer 架构，实现了低预算的扩散模型

介绍

论文地址：https://arxiv.org/abs/2407.15811
现代图像生成模型擅长创建自然、高质量的内容，每年生成的图像超过十亿幅。然而，从头开始训练这些模型极其昂贵和耗时。文本到图像（T2I）扩散模型降低了部分计算成本，但仍需要大量资源。

目前最先进的技术需要大约 18 000 个 A100 GPU 小时，而使用 8 个 H100 GPU 进行训练则需要一个多月的时间。此外，该技术通常依赖于大型或专有数据集，因此难以普及。

在这篇评论性论文中，我们开发了一种低成本、端到端文本到图像扩散建模管道，目的是在没有大型数据集的情况下显著降低成本。它侧重于基于视觉变换器的潜在扩散模型，利用其简单的设计和广泛的适用性。为了降低计算成本，可通过随机屏蔽输入标记来减少每幅图像需要处理的斑块数量。本文克服了现有遮蔽方法在高遮蔽率下性能下降的难题。

为了克服文本到图像扩散模型性能不佳的问题，本文提出了一种 "延迟掩蔽 "策略。通过在轻量级补丁混合器中处理补丁，然后将其输入扩散变换器，可以低成本实现可靠的训练，同时即使在高掩蔽率下也能保留语义信息。它还结合了变压器架构的最新发展，以提高大规模训练的性能。

该实验训练了一个 1.16 亿参数的稀疏扩散变换器，预算仅为 1,890 美元、3,700 万张图像和 75% 的屏蔽率。结果，在 COCO 数据集上零镜头生成的 FID 达到了 12.7。在一台 8×H100 GPU 机器上的训练时间仅为 2.6 天，与目前最先进的方法（37.6 天，GPU 成本 28,400 美元）相比缩短了 14 倍。

建议方法

延迟掩蔽

由于变换器的计算复杂度与序列的长度成正比，降低训练成本的一种方法是通过使用大尺寸补丁来减少序列，如图 1-b 所示。使用大尺寸补丁可使每幅图像的补丁数量呈二次方减少，但由于图像的大区域会被主动压缩成单个补丁，因此会显著降低性能。

一种方法是使用遮罩去除变换器输入层中的一些斑块，如图 1-c 所示，同时保持斑块大小。这种方法类似于卷积网络中的随机裁剪训练，但遮蔽补丁可以在图像的非连续区域进行训练。这种方法被广泛应用于视觉和语言领域。

图 1-d 中的MaskDiT还增加了补充自编码损失，以鼓励从遮蔽的斑块中学习表示法，从而促进遮蔽斑块的重建。这种方法屏蔽了 75% 的输入图像，大大降低了计算成本。

图 1.压缩补丁序列以降低计算成本。

然而，高遮罩率会大大降低转换器的整体性能：即使使用 MaskDiT，也只能看到与简单遮罩相比的微弱改进。这是因为即使采用这种方法，大部分图像斑块也会在输入层被去除。

本文引入了一个名为 "补丁混合器 "的预处理模块，用于在屏蔽之前处理补丁嵌入。这可确保未屏蔽的补丁保留整个图像的信息，从而提高学习效率。这种方法有可能提高性能，同时在计算上与现有的 MaskDiT 策略相当。

补丁混合器和学习障碍

补片混合器指的是任何能够融合单个补片嵌入的神经架构。在变换器模型中，这一目标自然可以通过注意机制和前馈层的结合来实现。因此，本文使用轻量级变压器（只有几层）作为补片混合器。输入序列标记经补丁混合器处理后，将被屏蔽（图 2e）。假定掩码为二进制 m，则使用以下损失函数对模型进行训练。

引入专家混合（MoE）和分层缩放的变换器架构

论文采用了先进变压器结构的创新技术，在计算受限的情况下提高了模型性能。

专家混合物（MoE，Zhou 等人，2022 年）：使用 MoE 层扩展模型的参数和表现力，同时避免训练成本的显著增加简化的 MoE 层与专家选择路由允许额外的辅助损失函数无需调整负载。
分层缩放 （Mehta 等人，2024 年 ）：这种方法已被证明能提高大型语言模型的性能，其中变换器块的宽度（隐藏层的维度）随深度线性增加。更多的参数被分配给更深的层，以学习更复杂的特征。

整体架构如图 2 所示。

图 2：拟议方法的总体概览。

试验

验证延迟遮蔽和补丁混合器的有效性

当许多补丁被遮蔽时，遮蔽性能会下降；Zheng 等人（2024 年）指出，当遮蔽率超过 50%，MaskDiT 的性能会显著下降。本文评估了遮蔽率高达 87.5% 时的性能，并将其与不使用补丁混合器的传统天真遮蔽方法进行了比较。本文中的 "延迟掩蔽 "使用了一个四层变压器块贴片混合器，其参数小于主干变压器参数的 10%。两者都使用了设置完全相同的 AdamW 优化器。

图 3 对结果进行了总结。延迟掩蔽在所有指标上都明显优于天真掩蔽和MaskDiT，表明随着掩蔽率的增加，性能差异也在扩大。例如，在屏蔽率为 75% 时，原始屏蔽的 FID 得分为 80，MaskDiT 的FID 得分为16.5，而拟议方法的 FID 得分为 5.03，优于未屏蔽时的 3.79。

图 3：验证延迟屏蔽和贴片混频器的有效性。

验证 "专家混合 "和 "分层缩放 "的有效性

分层缩放： 使用 DiT-Tiny 架构进行的实验比较了分层缩放和恒宽变换器与天真屏蔽。两个模型都在相同的计算负荷下进行了相同时间的训练。在所有性能指标上，逐层缩放方法始终优于恒定宽度模型，而且在屏蔽训练中更为有效。

专家混合物（MoE）： 测试了在交替区块中具有 MoE 层的 DiT-Tiny/2 变压器。总体性能与无 MoE 层的基线模型相似，Clip-score 略有提高（从 28.11 到 28.66），FID 分数有所下降（从 6.92 到 6.98）。改进幅度有限的原因是 60K 步的训练和每位专家看到的样本量较小。

与以往研究的比较

COCO 数据集（表 1）上的零镜头图像生成： 根据标题生成 30 000 幅图像，并使用 FID-30K 比较其与真实图像的分布情况。拟议方法的FID-30K 得分为 12.66，与之前的低成本训练方法相比，计算成本降低了 14 倍，而且不依赖于专有数据集。该方法的计算成本也比 Würstchen 低 19 倍（Pernias et al.）

表 1：在 COCO 数据集上生成零镜头图像

详细图像生成比较**（表 2）****：**GenEval（Ghosh 等人，2024 年）用于评估生成物体位置、共现、数量和颜色的能力。与 Stable-DiffusionXL-turbo 和 PixArt-α 模型相比，拟议方法在单个物体生成方面的准确性接近完美，与 Stable-Diffusion 变体相当，优于 Stable-Diffusion-1.5。在以下方面也表现出卓越的性能