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 秒出图（512×512)，且比 SD-v1.5 获得更好的 CLIP score。这是现在已知最快的端上 Stable Diffusion 模型！
-
论文地址：arxiv.org/pdf/2306.00…
-
Webpage: 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 等等）。
-
本文首要重视速度上的进步，并未对模型存储进行优化。咱们相信所提出的 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] www.qualcomm.com/news/onq/20…
[7] Speed Is All You Need: On-Device Acceleration of Large Diffusion Models via GPU-Aware Optimizations, CVPR’23 Workshop
