在手机等移动端侧运行 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 是如何优化得到的呢？

首先，让我们从问题出发，探讨为何优化是必要的

目前最热门的文本到图像生成技术都是基于扩散模型实现的。由于其预先训练的模型具有极强的基本图像生成能力和在下游微调任务上的稳健性质，因此我们看到了扩散模型在图像编辑、可控生成、个性化生成以及视频生成等领域的出色表现

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

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

模型优化

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

整体规划

在图上的左侧是原始 UNet 的设计示意，可以看出它基本上包括了卷积和Transformer，而Transformer又包括了自注意力机制和交叉注意力机制

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

经过实验证明，将整个UNet的16个Transformer移动到特征分辨率最低的内层，并且每一层都剪掉一个卷积，对性能没有明显影响。所达到的效果是：MobileDiffusion将原本的22个卷积和16个Transformer精简到了只有11个卷积和大约12个Transformer，并且这些注意力都是在低分辨率特征图上进行的。这样做的效率大大提高，带来了40%的效率提升和40%的参数剪切。最终的模型如右图所示。下面是与其他模型的对比：

需要进行重新编写的内容是：微观设计

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

解耦自注意力和交叉注意力

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

Finetune softmax into relu

众所周知，在大部分未优化的情况下，softmax函数很难进行并行处理，效率较低。MobileDiffusion提出了一种新的方法，即将softmax函数直接调整（finetune）为relu函数，因为relu函数对于每个数据点的激活更加高效。令人惊讶的是，仅需大约一万步的微调，模型的度量指标反而提高了，并且生成的图像质量也得到了保证。因此，相比于softmax函数，relu函数的优势显而易见

Separable Convolution (可分离卷积)

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

采样优化

目前最流行的采样优化方法包括渐进蒸馏（Progressive Distillation）和UFOGen，它们分别可以实现8步和1步。为了证明即使在模型经过极致简化后，这些采样方法仍然有效，MobileDiffusion对这两种方法进行了实验验证

优化后的采样与基准模型进行了比较，可以看出采样优化后的 8 步和 1 步模型的指标都有显著的提升

实验与应用

移动端基准测试

在 iPhone 15 Pro 上，MobileDiffusion 可以以当前最快的速度进行出图，只需 0.2 秒！

下游任务测试

MobileDiffusion 探索了包括 ControlNet/Plugin 和 LoRA Finetune 的下游任务。从下图可以看出，经过模型和采样优化后，MobileDiffusion 依然保持了优秀的模型微调能力。

总结

MobileDiffusion对多种模型和采样优化方法进行了探索，最终实现了在移动端亚秒级的图像生成能力，同时保证了下游微调应用的稳定性。我们相信这将对未来高效的扩散模型设计产生影响，并拓展移动端应用的实际应用案例