首頁 >科技週邊 >人工智慧 >iPhone兩秒出圖,目前已知最快的行動端Stable Diffusion模型來了

iPhone兩秒出圖,目前已知最快的行動端Stable Diffusion模型來了

王林
王林轉載
2023-06-12 11:15:481349瀏覽

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 模型!

iPhone两秒出图,目前已知的最快移动端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.25 x 加速),這是透過提出的CFG-aware Distillation 得到的。 VAE decoder 的加速是透過結構化剪枝來實現。

iPhone两秒出图,目前已知的最快移动端Stable Diffusion模型来了

以下著重於 Efficient UNet 的設計和 CFG-aware Distillation 損失函數的設計。

(1)Efficient UNet

我們透過分析UNet 中的Cross-Attention 和ResNet 模組,定位速度瓶頸在於Cross-Attention 模組(尤其是第一個Downsample 階段的Cross-Attention),如下圖所示。這個問題的根源是因為 attention 模組的複雜度跟特徵圖的 spatial size 成平方關係,在第一個 Downsample 階段,特徵圖的 spatial size 仍然較大,導致計算複雜度高。

iPhone两秒出图,目前已知的最快移动端Stable Diffusion模型来了

為了優化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 功能,目標是實現推理時間大幅縮減並維持影像品質不降,最能說明這一點的是下圖:

iPhone两秒出图,目前已知的最快移动端Stable Diffusion模型来了

該圖是在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 產生的樣本:

iPhone两秒出图,目前已知的最快移动端Stable Diffusion模型来了

更多樣本請參考文章附錄。

除了這些主要結果,文中也展示了許多燒蝕分析(Ablation Study)實驗,希望能為高效SD 模型的研發提供參考經驗:

#(1)之前Step Distillation 的工作通常採用漸進式方案[4, 5],但我們發現,在SD 模型上漸進式蒸餾並沒有比直接蒸餾更有優勢,且過程繁瑣,因此我們在文中採用的是直接蒸餾方案。

iPhone两秒出图,目前已知的最快移动端Stable Diffusion模型来了

(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 模型上更為有效。

iPhone两秒出图,目前已知的最快移动端Stable Diffusion模型来了

(3)既有Step Distillation 的工作[4, 5] 沒有將原有的損失函數和蒸餾損失函數加在一起,熟悉影像分類知識蒸餾的朋友應該知道,這種設計直覺上來說是欠優的。於是我們提出把原有的損失函數加入訓練中,如下圖所示,確實有效(小幅降低 FID)。

iPhone两秒出图,目前已知的最快移动端Stable Diffusion模型来了

總結與未來工作

本文提出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中文網其他相關文章!

陳述:
本文轉載於:51cto.com。如有侵權,請聯絡admin@php.cn刪除