機器之心專欄

機器之心編輯部

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

未來工作：

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

2. 本文主要關注速度上的提升，并未對模型存儲進行優化。我們相信所提出的 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 ] https://www.qualcomm.com/news/onq/2023/02/worlds-first-on-device-demonstration-of-stable-diffusion-on-android

[ 7 ] Speed Is All You Need: On-Device Acceleration of Large Diffusion Models via GPU-Aware Optimizations, CVPR ’ 23 Workshop

THE END

轉載請聯系本公眾號獲得授權

投稿或尋求報道：content@jiqizhixin.com