谷歌新作MobileDiffusion：移动设备上的亚秒级文本到图像生成

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

点击进入—>【扩散模型和Transformer】交流群

扫码加入CVer知识星球，可以最快学习到最新顶会顶刊上的论文idea和CV从入门到精通资料，以及最前沿项目和应用！发论文，强烈推荐！

转载自：机器之心

在手机等移动端侧运行 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 的模型，指标都是比较突出的。

实验与应用

移动端基准测试

MobileDiffusion 在 iPhone 15 Pro 上可以得到当前最快的出图速度，0.2s！

下游任务测试

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

总结

MobileDiffusion 探索了多种模型和采样优化方法，最终可以实现在移动端的亚秒级出图能力，下游微调应用依然有保障。我们相信这将会对今后高效的扩散模型设计产生影响，并拓展移动端应用实例。

CVPR / ICCV 2023论文和代码下载

后台回复：CVPR2023，即可下载CVPR 2023论文和代码开源的论文合集

后台回复：ICCV2023，即可下载ICCV 2023论文和代码开源的论文合集

扩散模型和Transformer交流群成立
扫描下方二维码，或者添加微信：CVer444，即可添加CVer小助手微信，便可申请加入CVer-扩散模型或者Transformer 微信交流群。另外其他垂直方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF等。
一定要备注：研究方向+地点+学校/公司+昵称（如扩散模型或者Transformer+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群

▲扫码或加微信号: CVer444，进交流群
CVer计算机视觉（知识星球）来了！想要了解最新最快最好的CV/DL/AI论文速递、优质实战项目、AI行业前沿、从入门到精通学习教程等资料，欢迎扫描下方二维码，加入CVer计算机视觉（知识星球），已汇集近万人！

▲扫码加入星球学习

▲点击上方卡片，关注CVer公众号

整理不易，请点赞和在看