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 倍。


