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

2023-12-03 05:12

机器之心专栏

机器之心编辑部

在手机等移动端侧运行 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

softmax 众所周知在大部分未优化情况下是难以并行的，因此效率很低。MobileDiffusion 提出直接将 softmax 函数 finetune 到 relu，因为 relu 是每一个点的激活，更为高效。令人惊讶的是，只需要大概一万步的微调，模型 metric 反而提升了，出图质量也有保障。因此 relu 相比于 softmax 的优势是很明显的了。

Separable Convolution (可分离卷积)

MobileDiffuison 精简参数的关键还在采用了 Seprable Convolution。这种技术已经被 MobileNet 等工作证实是极为有效的，特别是移动端，但是一般在生成模型上很少采用。MobileDiffusion 实验发现 Separable Convolution 对减少参数是很有效的，尤其是将其放在 UNet 最内层，模型质量经分析证明是没有损失的。

采样优化

当下最常采用的采样优化方法包括了 Progressive Distillation 和 UFOGen，分别可以做到 8 steps 和 1 step。为了证明在模型极致精简后，这些采样依然适用，MobileDiffusion 对两者同时做了实验验证。

采样优化前后和基准模型的比较如下，可以看出采样优化后的 8 steps 和 1 step 的模型，指标都是比较突出的。