Google发布设备端文生图模型MobileDiffusion

MobileDiffusion：设备端快速文生图

由杨赵（高级软件工程师）和侯廷波（高级员工软件工程师），Core ML 团队发布

文生图扩散模型在从文本提示生成高质量图像方面展现了卓越的能力。然而，业界领先的模型具有数十亿个参数，因此运行成本高昂，需要强大的台式机或服务器（例如 Stable Diffusion、DALL·E 和 Imagen）。虽然过去一年中通过 MediaPipe 在 Android 上和通过 Core ML 在 iOS 上的推理解决方案取得了进展，但在移动设备上实现快速（亚秒级）的文生图生成仍然遥不可及。

为此，我们在《MobileDiffusion：移动设备上的亚秒级文生图生成》中介绍了一种新颖方法，具有在设备端快速生成文生图的潜力。MobileDiffusion 是一个为移动设备量身定制的高效潜在扩散模型。我们还采用 DiffusionGAN 在推理过程中实现单步采样，该方法微调预训练的扩散模型，同时利用 GAN 来建模去噪步骤。我们在 iOS 和 Android 高端设备上测试了 MobileDiffusion，它能够在半秒内运行并生成 512x512 的高质量图像。其模型参数仅 520M 的相对较小规模使其特别适合移动设备部署。

文生图扩散模型效率相对较低源于两个主要挑战。首先，扩散模型的固有设计需要迭代去噪来生成图像，这要求多次评估模型。其次，文生图扩散模型中网络架构的复杂性涉及大量参数，通常达到数十亿级别，导致计算成本高昂。因此，尽管在移动设备上部署生成模型具有潜在优势，比如增强用户体验和应对新兴隐私问题，但这在当前文献中仍相对鲜少探讨。

文生图扩散模型推理效率的优化一直是一个活跃的研究领域。以往的研究主要集中在解决第一个挑战，即减少函数评估次数（NFEs）。通过利用先进的数值求解器（如 DPM）或蒸馏技术（如渐进蒸馏、一致性蒸馏），必要的采样步数已大幅从数百步减少到个位数。一些最近的技术，如 DiffusionGAN 和对抗扩散蒸馏，甚至能降低到单一步骤。

然而，在移动设备上，即使采样步数很少，由于模型架构的复杂性，速度也可能较慢。迄今为止，文生图扩散模型的架构效率相对关注较少。少数较早的工作简要涉及这一问题，涉及移除冗余神经网络块（如 SnapFusion）。然而，这些工作缺乏对模型架构内每个组件的综合分析，因此未能为设计高效架构提供全面指导。

要有效克服移动设备计算能力有限带来的挑战，需要对模型的架构效率进行深入而全面的探索。为了达到这一目标，我们的研究对 Stable Diffusion 的 UNet 架构中的每个组成部分和计算操作进行了详细检查。我们提出了一套全面指南来设计高效的文生图扩散模型，最终打造出 MobileDiffusion。

MobileDiffusion 的设计遵循潜在扩散模型的设计理念。它包含三个组件：文本编码器、扩散 UNet 和图像解码器。对于文本编码器，我们使用 CLIP-ViT/L14，这是一个适合移动设备的小型模型（125M 参数）。随后，我们将注意力转向扩散 UNet 和图像解码器。

如下图所示，扩散 UNets 通常交错使用 Transformer 块和卷积块。我们对这两个基础构件进行了综合调研。在研究过程中，我们控制训练管道（如数据、优化器）以研究不同架构的效果。

在经典的文生图扩散模型中，一个 Transformer 块包括自注意力层（SA）用于建模视觉特征的长距离依赖、交叉注意力层（CA）以捕捉文本条件与视觉特征之间的交互，以及前馈层（FF）来后处理注意力层的输出。这些 Transformer 块在文生图扩散模型中发挥关键作用，是负责文本理解的主要组件。然而，由于注意力操作的计算复杂度与序列长度呈二次方关系，它们也构成了显著的效率挑战。我们遵循 UViT 架构的思想，在 UNet 的瓶颈处放置更多的 Transformer 块。这一设计选择的动机在于注意力计算在瓶颈处的资源消耗较少，因为其维度更低。

卷积块，特别是 ResNet 块，部署在 UNet 的每个级别。虽然这些块对特征提取和信息流动至关重要，但相关的计算成本，尤其是在高分辨率级别，可能相当可观。一个已被证实的方法是可分离卷积。我们发现在 UNet 的更深层部分用轻量级可分离卷积层替换常规卷积层能获得相似的性能。

在下图中，我们比较了几个扩散模型的 UNets。我们的 MobileDiffusion 在 FLOPs（浮点运算）和参数数量方面表现出卓越的效率。

除了 UNet 之外，我们还优化了图像解码器。我们训练了一个变分自编码器（VAE）将 RGB 图像编码为 8 通道潜在变量，图像的空间大小缩小 8 倍。潜在变量可以解码为图像，并放大 8 倍。为进一步提升效率，我们通过剪枝原始解码器的宽度和深度来设计轻量级解码器架构。由此产生的轻量级解码器带来了显著的性能提升，延迟改善近 50%，质量也更好。更多细节请参考我们的论文。

除了优化模型架构外，我们还采用了 DiffusionGAN 混合方法来实现单步采样。为文生图生成训练 DiffusionGAN 混合模型存在多个复杂之处。值得注意的是，判别器作为区分真实数据和生成数据的分类器，必须基于纹理和语义进行判断。此外，训练文生图模型的成本可能极其高昂，特别是在基于 GAN 的模型中，判别器引入了额外参数。纯 GAN 的文生图模型（如 StyleGAN-T、GigaGAN）面临类似的复杂性，导致训练高度复杂且成本高昂。

为克服这些挑战，我们使用预训练的扩散 UNet 来初始化生成器和判别器。这一设计使得能够与预训练扩散模型无缝初始化。我们推测扩散模型内的内部特征包含有关文本和视觉数据之间复杂相互作用的丰富信息。这一初始化策略显著简化了训练过程。

下图展示了训练流程。初始化后，噪声图像被发送到生成器进行单步扩散。其结果与真值通过重构损失进行评估，类似于扩散模型训练。我们随后向输出添加噪声，并将其发送到判别器，其结果通过 GAN 损失进行评估，实际上采用 GAN 来建模去噪步骤。通过使用预训练权重初始化生成器和判别器，训练成为一个微调过程，可在少于 10K 次迭代内收敛。

下面我们展示了使用 DiffusionGAN 单步采样的 MobileDiffusion 生成的示例图像。凭借如此紧凑的模型（总共 520M 参数），MobileDiffusion 能够为各种域生成高质量多样化的图像。

我们在 iOS 和 Android 设备上测试了 MobileDiffusion 的性能，使用了不同的运行时优化器。延迟数字报告如下。我们发现 MobileDiffusion 效率非常高，能够在半秒内运行生成 512x512 的图像。这种闪电般的速度可能为移动设备上启用许多有趣的用例。

凭借在延迟和模型大小方面的优越效率，考虑到其能够在用户输入文本提示时实现快速图像生成体验的能力，MobileDiffusion 有潜力成为移动部署的一个非常友好的选择。我们将确保这项技术的任何应用都符合 Google 的负责任 AI 实践。

我们感谢所有帮助将 MobileDiffusion 带到设备端的合作者和贡献者：肖志生、许燕武、唐久强、贾浩琳、Lutz Justen、Daniel Fenner、Ronald Wotzlaw、韦佳宁、Raman Sarokin、Juhyun Lee、Andrei Kulik、Chang Chuo-Ling 和 Matthias Grundmann。