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 模型！
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.25 x 加速），這是透過提出的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.4 x 加速；（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 更進一步。
以上是iPhone兩秒出圖，目前已知最快的行動端Stable Diffusion模型來了的詳細內容。更多資訊請關注PHP中文網其他相關文章！
