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 加速?？傮w來說，SnapFusion 在 iPhone 14 Pro 上實現 2 秒內出圖，這是目前已知最快的移動端 Stable Diffusion 模型。

未來工作：

1.SD 模型在多種圖像生成場景中都可以使用，本文囿于時間，目前只關注了 text to image 這個核心任務，后期將跟進其他任務（如 inpainting，ControlNet 等等）。

2. 本文主要關注速度上的提升，并未對模型存儲進行優化。我們相信所提出的 Efficient UNet 仍然具備壓縮的空間，結合其他的高性能優化方法（如剪枝，量化），有望縮小存儲，并將時間降低到 1 秒以內，離端上實時 SD 更進一步。