MobileDiffusion：设备上快速生成文本到图像

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

为此，我们在“ MobileDiffusion：移动设备上的亚秒级文本转图像生成”中介绍了一种新方法，该方法有可能在设备上快速生成文本转图像。MobileDiffusion 是一种专为移动设备设计的高效潜在扩散模型。我们还采用DiffusionGAN在推理过程中实现一步采样，该模型微调预先训练的扩散模型，同时利用 GAN 来建模去噪步骤。我们已经在 iOS 和 Android 高级设备上测试了 MobileDiffusion，它可以在半秒内运行并生成 512x512 的高质量图像。其模型大小相对较小，只有 520M 个参数，非常适合移动部署。

在设备上快速生成文本到图像。

背景

文本到图像扩散模型相对低效的原因有两个主要方面。首先，扩散模型的固有设计要求迭代去噪来生成图像，因此需要对模型进行多次评估。其次，文本到图像扩散模型的网络架构复杂，涉及大量参数，通常达到数十亿个，因此需要进行计算成本高昂的评估。因此，尽管在移动设备上部署生成模型具有潜在优势，例如增强用户体验和解决新出现的隐私问题，但在当前文献中，该领域仍相对未被深入探索。

文本到图像扩散模型中推理效率的优化一直是一个活跃的研究领域。先前的研究主要集中在解决第一个挑战，寻求减少函数评估 (NFE) 的数量。利用先进的数值求解器（例如DPM）或蒸馏技术（例如渐进式蒸馏、一致性蒸馏），必要的采样步骤数量已从几百个显著减少到个位数。一些最近的技术，如DiffusionGAN和对抗性扩散蒸馏，甚至减少到一个必要的步骤。

然而，在移动设备上，由于模型架构的复杂性，即使是少量的评估步骤也可能很慢。到目前为止，文本到图像扩散模型的架构效率受到的关注相对较少。一些早期的研究简要涉及这个问题，涉及删除冗余的神经网络块（例如，SnapFusion）。然而，这些努力缺乏对模型架构中每个组件的全面分析，因此无法为设计高效的架构提供整体指南。

移动传播

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

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

扩散 UNet

如下图所示，扩散 UNets 通常交错变压器块和卷积块。我们对这两个基本构建块进行了全面调查。在整个研究过程中，我们控制训练管道（例如数据、优化器）以研究不同架构的效果。

在经典的文本到图像扩散模型中，Transformer 块由一个自注意层 (SA)（用于建模视觉特征之间的长程依赖关系）、一个交叉注意层 (CA)（用于捕获文本条件和视觉特征之间的交互）和一个前馈层 (FF)（用于后处理注意层的输出）组成。这些 Transformer 块在文本到图像扩散模型中起着关键作用，是负责文本理解的主要组件。然而，考虑到注意操作的计算成本（该成本与序列长度成二次方），它们也带来了巨大的效率挑战。我们遵循UViT架构的思想，该架构在 UNet 的瓶颈处放置更多的 Transformer 块。这种设计选择的动机是由于注意力计算在瓶颈处由于维度较低而消耗的资源较少。

我们的 UNet 架构在中间包含了更多的变换器，并在更高分辨率下跳过了自注意力 (SA) 层。

卷积块（特别是ResNet块）部署在 UNet 的每个级别。虽然这些块对于特征提取和信息流至关重要，但相关的计算成本（尤其是在高分辨率级别）可能非常高。在这种情况下，一种行之有效的方法是可分离卷积。我们观察到，在 UNet 的较深部分中，用轻量级可分离卷积层替换常规卷积层可获得类似的性能。

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

一些扩散 UNets 的比较。

图像解码器

除了 UNet，我们还优化了图像解码器。我们训练了一个变分自动编码器(VAE)，将RGB图像编码为 8 通道潜在变量，图像的空间大小减小了 8 倍。潜在变量可以解码为图像，大小增大 8 倍。为了进一步提高效率，我们通过修剪原始图像的宽度和深度设计了一个轻量级解码器架构。由此产生的轻量级解码器可显著提高性能，延迟改善近 50%，质量也更好。有关更多详细信息，请参阅我们的论文。

VAE 重建。我们的 VAE 解码器比 SD（稳定扩散）具有更好的视觉质量。

解码器   #参数（M）    峰值信噪比↑    SSIM↑    低功率IPS↓  

标准差 49.5 26.7 0.76 0.037

我们的 39.3 30.0 0.83 0.032

我们的-Lite    9.8 30.2 0.84 0.032

VAE 解码器的质量评估。我们的精简解码器比 SD 小得多，具有更好的质量指标，包括峰值信噪比(PSNR)、结构相似性指数测量(SSIM) 和学习感知图像块相似性(LPIPS)。

一步取样

除了优化模型架构外，我们还采用了DiffusionGAN 混合模型来实现一步采样。训练 DiffusionGAN 混合模型进行文本转图像生成会遇到一些复杂问题。值得注意的是，鉴别器（区分真实数据和生成数据的分类器）必须根据纹理和语义做出判断。此外，训练文本转图像模型的成本可能非常高，特别是在基于 GAN 的模型中，鉴别器会引入额外的参数。纯基于 GAN 的文本转图像模型（例如StyleGAN-T、GigaGAN）也面临着类似的复杂性，导致训练非常复杂且成本高昂。

为了克服这些挑战，我们使用预训练的扩散 UNet 来初始化生成器和鉴别器。这种设计可以无缝初始化预训练的扩散模型。我们假设扩散模型中的内部特征包含文本和视觉数据之间复杂相互作用的丰富信息。这种初始化策略大大简化了训练。

下图说明了训练过程。初始化后，将噪声图像发送到生成器进行一步扩散。使用重建损失对结果与真实值进行比较，类似于扩散模型训练。然后，我们将噪声添加到输出并将其发送到鉴别器，使用 GAN 损失对鉴别器的结果进行评估，有效地采用 GAN 来模拟去噪步骤。通过使用预训练权重来初始化生成器和鉴别器，训练变成了一个微调过程，在不到 10K 次迭代中收敛。

DiffusionGAN 微调的图示。

结果

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

图片由我们的 MobileDiffusion 生成

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

移动设备上的延迟测量（s ）。

结论

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

致谢

我们要感谢帮助将 MobileDiffusion 引入设备端的合作者和贡献者：Zhisheng Xiao、Yanwu Xu、Jiuqiang Tang、Haolin Jia、Lutz Justen、Daniel Fenner、Ronald Wotzlaw、Jianing Wei、Raman Sarokin、Juhyun Lee、Andrei Kulik、Chuo-Ling Chang 和 Matthias Grundmann。