手机上0.2秒出图、当前速度之最，谷歌打造超快扩散模型MobileDiffusion

原文来源：机器之心

图片来源：由无界 AI生成

在手机等移动端侧运行 Stable Diffusion 等文生图生成式 AI 大模型已经成为业界追逐的热点之一，其中生成速度是主要的制约因素。

近日，来自谷歌的一篇论文「MobileDiffusion: Subsecond Text-to-Image Generation on Mobile Devices」，提出了手机端最快文生图，在 iPhone 15 Pro 上只要 0.2 秒。论文出自 UFOGen 同一团队，在打造超小扩散模型的同时， 采用当前大火的 Diffusion GAN 技术路线做采样加速。

论文地址：https://arxiv.org/abs/2311.16567

下面是 MobileDiffusion 一步生成的结果。

那么，MobileDiffusion 是如何优化得到的呢？

让我们先从问题出发，为什么优化是必须的。

当下最火的文本到图像生成都是基于扩散模型来实现的。依赖于其预训练模型强大的基本图像生成能力和在下游微调任务上的稳健性质， 我们看到了扩散模型在诸如图像编辑、可控生成、 个性化生成以及视频生成的非凡表现。

然而作为 Foundation Model， 它的不足也很明显，主要包括了两方面：一是扩散模型的大量参数导致计算速度慢，尤其是在资源有限的情况下；二是扩散模型需要多步才能采样，这进一步导致很慢的推理速度。拿最火的的 Stable Diffusion 1.5 (SD) 为例，其基础模型包含了接近 10 亿参数，我们在 iPhone 15 Pro 上将模型量化后进行推理，50 步采样需要接近 80s。如此昂贵的资源需求和迟滞的用户体验极大的限制了其在移动端的应用场景。

为了解决以上问题，MobileDiffusion 点对点地进行优化。(1) 针对模型体积庞大的问题，我们主要对其核心组件 UNet 进行了大量试验及优化，包括了将计算昂贵的卷积精简和注意力运算放在了较低的层上，以及针对 Mobile Devices 的操作优化，诸如激活函数等。（2）针对扩散模型需要多步采样的问题， MobileDiffusion 探索并实践了像 Progressive Distillation 和当前最先进的 UFOGen 的一步推理技术。

模型优化

MobileDiffusion 基于当下开源社区里最火的 SD 1.5 UNet 进行优化。在每次的优化操作后， 会同时衡量相对于原始 UNet 模型的性能的损失，测量指标包括 FID 和 CLIP 两个常用 metric。

宏观设计

上图左边是原始 UNet 的设计示意， 可以看出基本包括了 Convolution 和 Transformer，Transformer 又包括了 Self-Attention 和 Cross-Attention。

MobileDiffusion 对 UNet 优化的核心思路分为两点：1）精简 Convolution， 众所周知，在高分辨率的特征空间上进行了 Convolution 是十分耗时的， 而且参数量很大，这里指的是 Full Convolution；2）提高 Attention 效率。和 Convolution 一样，高 Attention 需要对整个特征空间的长度进行运算，Self-Attention 复杂度和特征空间展平后长度成平方关系，Cross-Attention 也要和空间长度成正比。

实验表明将整个 UNet 的 16 个 Transformer 移到特征分辨率最低的内层，同时每一层都剪掉一个卷积，不会对性能有明显影响。达到的效果就是：MobileDiffusion 将原本 22 个 Convolution 和 16 个 Transformer，可以极致地精简到 11 个 Convolution 和 12 个左右 Transformer，并且这些注意力都是在低分辨率特征图上进行的，因为效率会极大提升，带来了 40% 效率提升和 40% 参数剪切，最终模型如上图右图所示。和更多模型的对比如下：

微观设计

这里将只介绍几种新颖的设计，有兴趣的读者可以阅读正文， 会有更详细的介绍。

Decouple Self-Attention and Cross-Attention

传统 UNet 里 Transformer 同时包含 Self-Attention 和 Cross-Attention，MobileDiffusion 将 Self-Attention 全部放在了最低分辨率特征图，但是保留一个 Cross-Attention 在中间层，发现这种设计既提高了运算效率又保证了模型出图质量

Finetune softmax into relu