Stable Diffusion（SD）是当前最热门的文本到图像（text to image）生成扩散模型。尽管其强大的图像生成能力令人震撼，一个明显的不足是需要的计算资源巨大，推理速度很慢：以 SD-v1.5 为例，即使用半精度存储，其模型大小也有 1.7GB，近 10 亿参数，端上推理时间往往要接近 2min。
为了解决推理速度问题，学术界与业界已经开始对 SD 加速的研究，主要集中于两条路线：
- 减少推理步数，这条路线又可以分为两条子路线，一是通过提出更好的 noise scheduler 来减少步数，代表作是 DDIM、PNDM、DPM 等；二是通过渐进式蒸馏（Progressive Distillation）来减少步数，代表作是 Progressive Distillation 和 w-conditioning 等。
- 工程技巧优化，代表作是 Qualcomm 通过 int8 量化 + 全栈式优化实现 SD-v1.5 在安卓手机上 15s 出图，Google 通过端上 GPU 优化将 SD-v1.4 在三星手机上加速到 12s。
尽管这些工作取得了长足的进步，但仍然不够快。
近日，Snap 研究院推出最新高性能 Stable Diffusion 模型，通过对网络结构、训练流程、损失函数全方位进行优化，在 iPhone 14 Pro 上实现 2 秒出图（512x512)，且比 SD-v1.5 取得更好的 CLIP score。这是目前已知最快的端上 Stable Diffusion 模型！
核心方法
Stable Diffusion 模型分为三部分：VAE encoder/decoder, text encoder, UNet，其中 UNet 无论是参数量还是计算量，都占绝对的大头，因此 SnapFusion 主要是对 UNet 进行优化。具体分为两部分：
- UNet 结构上的优化：通过分析原有 UNet 的速度瓶颈，本文提出一套 UNet 结构自动评估、进化流程，得到了更为高效的 UNet 结构（称为 Efficient UNet）。
- 推理步数上的优化：众所周知，扩散模型在推理时是一个迭代的去噪过程，迭代的步数越多，生成图片的质量越高，但时间代价也随着迭代步数线性增加。为了减少步数并维持图片质量，我们提出一种 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）来加速，但是它们没有在蒸馏训练中把 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 倍。
