iPhone两秒出图，目前已知的最快移动端Stable Diffusion模型来了
机器之心专栏
机器之心编辑部
Stable Diffusion （SD）是当前最热门的文本到图像（text to image）生成扩散模型。尽管其强大的图像生成能力令人震撼，一个明显的不足是需要的计算资源巨大，推理速度很慢：以 SD-v1.5 为例，即使用半精度存储，其模型大小也有 1.7GB，近 10 亿参数，端上推理时间往往要接近 2min。
为了解决推理速度问题，学术界与业界已经开始对 SD 加速的研究，主要集中于两条路线：（1）减少推理步数，这条路线又可以分为两条子路线，一是通过提出更好的 noise scheduler 来减少步数，代表作是 DDIM [1]，PNDM [2]，DPM [3] 等；二是通过渐进式蒸馏（Progressive Distillation）来减少步数，代表作是 Progressive Distillation [4] 和 w-conditioning [5] 等。（2）工程技巧优化，代表作是 Qualcomm 通过 int8 量化 + 全栈式优化实现 SD-v1.5 在安卓手机上 15s 出图 [6]，Google 通过端上 GPU 优化将 SD-v1.4 在三星手机上加速到 12s [7]。
尽管这些工作取得了长足的进步，但仍然不够快。
近日，Snap 研究院推出最新高性能 Stable Diffusion 模型，通过对网络结构、训练流程、损失函数全方位进行优化，在 iPhone 14 Pro 上实现 2 秒出图（512x512)，且比 SD-v1.5 取得更好的 CLIP score。这是目前已知最快的端上 Stable Diffusion 模型！
- 论文地址：https://arxiv.org/pdf/2306.00980.pdf
- Webpage: https://snap-research.github.io/SnapFusion
核心方法
Stable Diffusion 模型分为三部分：VAE encoder/decoder, text encoder, UNet，其中 UNet 无论是参数量还是计算量，都占绝对的大头，因此 SnapFusion 主要是对 UNet 进行优化。具体分为两部分：（1）UNet 结构上的优化：通过分析原有 UNet 的速度瓶颈，本文提出一套 UNet 结构自动评估、进化流程，得到了更为高效的 UNet 结构（称为 Efficient UNet）。（2）推理步数上的优化：众所周知，扩散模型在推理时是一个迭代的去噪过程，迭代的步数越多，生成图片的质量越高，但时间代价也随着迭代步数线性增加。为了减少步数并维持图片质量，我们提出一种 CFG-aware 蒸馏损失函数，在训练过程中显式考虑 CFG （Classifier-Free Guidance）的作用，这一损失函数被证明是提升 CLIP score 的关键！
下表是 SD-v1.5 与 SnapFusion 模型的概况对比，可见速度提升来源于 UNet 和 VAE decoder 两个部分，UNet 部分是大头。UNet 部分的改进有两方面，一是单次 latency 下降（1700ms -> 230ms，7.4x 加速），这是通过提出的 Efficient UNet 结构得到的；二是 Inference steps 降低（50 -> 8，6.25x 加速），这是通过提出的 CFG-aware Distillation 得到的。VAE decoder 的加速是通过结构化剪枝实现。
下面着重介绍 Efficient UNet 的设计和 CFG-aware Distillation 损失函数的设计。
（1）Efficient UNet
我们通过分析 UNet 中的 Cross-Attention 和 ResNet 模块，定位速度瓶颈在于 Cross-Attention 模块（尤其是第一个 Downsample 阶段的 Cross-Attention），如下图所示。这个问题的根源是因为 attention 模块的复杂度跟特征图的 spatial size 成平方关系，在第一个 Downsample 阶段，特征图的 spatial size 仍然较大，导致计算复杂度高。
为了优化 UNet 结构，我们提出一套 UNet 结构自动评估、进化流程：先对 UNet 进行鲁棒性训练（Robust Training），在训练中随机 drop 一些模块，以此来测试出每个模块对性能的真实影响，从而构建一个 “对 CLIP score 的影响 vs. latency” 的查找表；然后根据该查找表，优先去除对 CLIP score 影响不大同时又很耗时的模块。这一套流程是在线自动进行，完成之后，我们就得到了一个全新的 UNet 结构，称为 Efficient UNet。相比原版 UNet，实现 7.4x 加速且性能不降。
（2）CFG-aware Step Distillation
CFG（Classifier-Free Guidance）是 SD 推理阶段的必备技巧，可以大幅提升图片质量，非常关键！尽管已有工作对扩散模型进行步数蒸馏（Step Distillation）来加速 [4]，但是它们没有在蒸馏训练中把 CFG 纳入优化目标，也就是说，蒸馏损失函数并不知道后面会用到 CFG。这一点根据我们的观察，在步数少的时候会严重影响 CLIP score。
为了解决这个问题，我们提出在计算蒸馏损失函数之前，先让 teacher 和 student 模型都进行 CFG，这样损失函数是在经过 CFG 之后的特征上计算，从而显式地考虑了不同 CFG scale 的影响。实验中我们发现，完全使用 CFG-aware Distillation 尽管可以提高 CLIP score， 但 FID 也明显变差。我们进而提出了一个随机采样方案来混合原来的 Step Distillation 损失函数和 CFG-aware Distillation 损失函数，实现了二者的优势共存，既显著提高了 CLIP score，同时 FID 也没有变差。这一步骤，实现进一步推理阶段加速 6.25 倍，实现总加速约 46 倍。
除了以上两个主要贡献，文中还有对 VAE decoder 的剪枝加速以及蒸馏流程上的精心设计，具体内容请参考论文。
实验结果
SnapFusion 对标 SD-v1.5 text to image 功能，目标是实现推理时间大幅缩减并维持图像质量不降，最能说明这一点的是下图：
该图是在 MS COCO’14 验证集上随机选取 30K caption-image pairs 测算 CLIP score 和 FID。CLIP score 衡量图片与文本的语义吻合程度，越大越好；FID 衡量生成图片与真实图片之间的分布距离（一般被认为是生成图片多样性的度量），越小越好。图中不同的点是使用不同的 CFG scale 获得，每一个 CFG scale 对应一个数据点。从图中可见，我们的方法（红线）可以达到跟 SD-v1.5（蓝线）同样的最低 FID，同时，我们方法的 CLIP score 更好。值得注意的是，SD-v1.5 需要 1.4min 生成一张图片，而 SnapFusion 仅需要 1.84s，这也是目前我们已知最快的移动端 Stable Diffusion 模型！
下面是一些 SnapFusion 生成的样本：
更多样本请参考文章附录。
除了这些主要结果，文中也展示了众多烧蚀分析（Ablation Study）实验，希望能为高效 SD 模型的研发提供参考经验：
（1）之前 Step Distillation 的工作通常采用渐进式方案 [4, 5]，但我们发现，在 SD 模型上渐进式蒸馏并没有比直接蒸馏更有优势，且过程繁琐，因此我们在文中采用的是直接蒸馏方案。
（2）CFG 虽然可以大幅提升图像质量，但代价是推理成本翻倍。今年 CVPR’23 Award Candidate 的 On Distillation 一文 [5] 提出 w-conditioning，将 CFG 参数作为 UNet 的输入进行蒸馏（得到的模型叫做 w-conditioned UNet），从而在推理时省却 CFG 这一步，实现推理成本减半。但是我们发现，这样做其实会造成图片质量下降，CLIP score 降低（如下图中，四条 w-conditioned 线 CLIP score 均未超过 0.30, 劣于 SD-v1.5）。而我们的方法则可以减少步数，同时将 CLIP score 提高，得益于所提出的 CFG-aware 蒸馏损失函数！尤其值得主要的是，下图中绿线（w-conditioned, 16 steps）与橙线（Ours，8 steps）的推理代价是一样的，但明显橙线更优，说明我们的技术路线比 w-conditioning [5] 在蒸馏 CFG guided SD 模型上更为有效。
（3）既有 Step Distillation 的工作 [4, 5] 没有将原有的损失函数和蒸馏损失函数加在一起，熟悉图像分类知识蒸馏的朋友应该知道，这种设计直觉上来说是欠优的。于是我们提出把原有的损失函数加入到训练中，如下图所示，确实有效（小幅降低 FID）。
总结与未来工作
本文提出 SnapFusion，一种移动端高性能 Stable Diffusion 模型。SnapFusion 有两点核心贡献：（1）通过对现有 UNet 的逐层分析，定位速度瓶颈，提出一种新的高效 UNet 结构（Efficient UNet），可以等效替换原 Stable Diffusion 中的 UNet，实现 7.4x 加速；（2）对推理阶段的迭代步数进行优化，提出一种全新的步数蒸馏方案（CFG-aware Step Distillation），减少步数的同时可显著提升 CLIP score，实现 6.25x 加速。总体来说，SnapFusion 在 iPhone 14 Pro 上实现 2 秒内出图，这是目前已知最快的移动端 Stable Diffusion 模型。
未来工作：
1.SD 模型在多种图像生成场景中都可以使用，本文囿于时间，目前只关注了 text to image 这个核心任务，后期将跟进其他任务（如 inpainting，ControlNet 等等）。
2. 本文主要关注速度上的提升，并未对模型存储进行优化。我们相信所提出的 Efficient UNet 仍然具备压缩的空间，结合其他的高性能优化方法（如剪枝，量化），有望缩小存储，并将时间降低到 1 秒以内，离端上实时 SD 更进一步。
参考文献
[1] Denoising Diffusion Implicit Models, ICLR’21
[2] Pseudo Numerical Methods for Diffusion Models on Manifolds, ICLR’22
[3] DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps, NeurIPS’22
[4] Progressive Distillation for Fast Sampling of Diffusion Models, ICLR’22
[5] On Distillation of Guided Diffusion Models, CVPR’23
[6] https://www.qualcomm.com/news/onq/2023/02/worlds-first-on-device-demonstration-of-stable-diffusion-on-android
[7] Speed Is All You Need: On-Device Acceleration of Large Diffusion Models via GPU-Aware Optimizations, CVPR’23 Workshop