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

2024年1月31日 | 作者：Yang Zhao（高级软件工程师）、Tingbo Hou（高级主任软件工程师，Core ML团队）

概述

文本到图像扩散模型在根据文本提示生成高质量图像方面展现出卓越能力。然而，主流模型包含数十亿参数，运行成本高昂，需要强大台式机或服务器支持（例如Stable Diffusion、DALL·E和Imagen）。尽管过去一年通过MediaPipe在Android和Core ML在iOS上的推理解决方案取得进展，但在移动设备上实现快速（亚秒级）文本到图像生成仍然难以实现。

为此，我们在《MobileDiffusion：移动设备上的亚秒级文本到图像生成》中提出了一种新颖方法，具有在设备上快速生成文本到图像的潜力。MobileDiffusion是专为移动设备设计的高效潜在扩散模型。我们采用DiffusionGAN在推理过程中实现单步采样，通过微调预训练扩散模型并利用GAN对去噪步骤进行建模。我们在iOS和Android高端设备上测试了MobileDiffusion，可在半秒内生成512x512高质量图像。其相对较小的模型尺寸（仅5.2亿参数）使其特别适合移动端部署。

技术背景

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

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

然而，在移动设备上，即使少量评估步骤也可能因模型架构的复杂性而变得缓慢。迄今为止，文本到图像扩散模型的架构效率受到的关注相对较少。早期的一些工作简要涉及此问题，包括移除冗余神经网络块（例如SnapFusion）。但这些努力缺乏对模型架构中每个组件的全面分析，因此未能为设计高效架构提供全面指导。

MobileDiffusion技术架构

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

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

扩散UNet架构

如下图所示，扩散UNet通常交替使用Transformer块和卷积块。我们对这两种基本构建块进行了全面研究。在整个研究中，我们控制训练流程（例如数据、优化器）以研究不同架构的效果。

在经典文本到图像扩散模型中，Transformer块包含自注意力层（SA）用于建模视觉特征间的长程依赖关系，交叉注意力层（CA）用于捕获文本条件与视觉特征间的交互，以及前馈层（FF）用于后处理注意力层的输出。这些Transformer块在文本到图像扩散模型中起着关键作用，是负责文本理解的主要组件。然而，它们也带来显著的效率挑战，因为注意力操作的计算成本与序列长度呈二次方关系。

我们遵循UViT架构的思想，在UNet的瓶颈处放置更多Transformer块。这种设计选择的动机是：在瓶颈处由于维度较低，注意力计算资源密集度较低。我们的UNet架构在中间部分包含更多Transformer，并在更高分辨率处跳过自注意力（SA）层。

卷积块（特别是ResNet块）部署在UNet的每个级别。虽然这些块有助于特征提取和信息流，但相关的计算成本（尤其是在高分辨率级别）可能相当大。在这方面，一种经过验证的方法是分离卷积。我们观察到，在UNet的较深部分用轻量级分离卷积层替换常规卷积层可产生相似性能。

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

图像解码器优化

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

VAE重建质量对比显示，我们的VAE解码器比SD（Stable Diffusion）具有更好的视觉质量。

解码器性能评估表：

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

单步采样技术

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

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

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

性能结果

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

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

移动设备延迟测量（秒）数据表展示了在不同设备和优化器上的性能表现。

结论

凭借在延迟和尺寸方面的卓越效率，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。