機器之心專欄
機器之心編輯部
在手機等移動端側運行 Stable Diffusion 等文生圖生成式 AI 大模型已經成為業界追逐的熱點之一，其中生成速度是主要的制約因素。
近日，來自谷歌的一篇論文「MobileDiffusion: Subsecond Text-to-Image Generation on Mobile Devices」，提出了手機端最快文生圖，在 iPhone 15Pro 上只要 0.2 秒。論文出自 UFOGen 同一團隊，在打造超小擴散模型的同時， 採用當前大火的 Diffusion GAN 技術路線做採樣加速。
論文地址：https://arxiv.org/abs/2311.16567
下面是 MobileDiffusion 一步生成的結果。
那麼，MobileDiffusion 是如何優化得到的呢？
讓我們先從問題出發，為什麼優化是必須的。
當下最火的文本到圖像生成都是基於擴散模型來實現的。依賴於其預訓練模型強大的基本圖像生成能力和在下游微調任務上的穩健性質， 我們看到了擴散模型在諸如圖像編輯、可控生成、 個性化生成以及視頻生成的非凡表現。
然而作為 Foundation Model， 它的不足也很明顯，主要包括了兩方面：一是擴散模型的大量參數導致計算速度慢，尤其是在資源有限的情況下；二是擴散模型需要多步才能採樣，這進一步導致很慢的推理速度。拿最火的的 Stable Diffusion 1.5 ( SD ) 為例，其基礎模型包含了接近 10 億參數，我們在 iPhone 15 Pro 上將模型量化後進行推理，50 步採樣需要接近 80s。如此昂貴的資源需求和遲滯的用戶體驗極大的限制了其在移動端的應用場景。
為了解決以上問題，MobileDiffusion 點對點地進行優化。 ( 1 ) 針對模型體積龐大的問題，我們主要對其核心組件 UNet 進行了大量試驗及優化，包括了將計算昂貴的卷積精簡和注意力運算放在了較低的層上，以及針對 Mobile Devices 的操作優化，諸如激活函數等。（2）針對擴散模型需要多步採樣的問題， MobileDiffusion 探索並實踐了像 Progressive Distillation 和當前最先進的 UFOGen 的一步推理技術。
模型優化
MobileDiffusion 基於當下開源社區里最火的 SD 1.5 UNet 進行優化。在每次的優化操作後， 會同時衡量相對於原始 UNet 模型的性能的損失，測量指標包括 FID 和 CLIP 兩個常用 metric。
宏觀設計
上圖左邊是原始 UNet 的設計示意， 可以看出基本包括了 Convolution 和 Transformer，Transformer 又包括了 Self-Attention 和 Cross-Attention。
MobileDiffusion 對 UNet 優化的核心思路分為兩點：1）精簡 Convolution， 眾所周知，在高分辨率的特徵空間上進行了 Convolution 是十分耗時的， 而且參數量很大，這裡指的是 Full Convolution；2）提高 Attention 效率。和 Convolution 一樣，高 Attention 需要對整個特徵空間的長度進行運算，Self-Attention 複雜度和特徵空間展平後長度成平方關係，Cross-Attention 也要和空間長度成正比。
實驗表明將整個 UNet 的 16 個 Transformer 移到特徵分辨率最低的內層，同時每一層都剪掉一個卷積，不會對性能有明顯影響。達到的效果就是：MobileDiffusion 將原本 22 個 Convolution 和 16 個 Transformer，可以極致地精簡到 11 個 Convolution 和 12 個左右 Transformer，並且這些注意力都是在低分辨率特徵圖上進行的，因為效率會極大提升，帶來了 40% 效率提升和 40% 參數剪切，最終模型如上圖右圖所示。和更多模型的對比如下：
微觀設計
這裡將只介紹幾種新穎的設計，有興趣的讀者可以閱讀正文， 會有更詳細的介紹。
傳統 UNet 里 Transformer 同時包含 Self-Attention 和 Cross-Attention，MobileDiffusion 將 Self-Attention 全部放在了最低分辨率特徵圖，但是保留一個 Cross-Attention 在中間層，發現這種設計既提高了運算效率又保證了模型出圖質量
softmax 眾所周知在大部分未優化情況下是難以並行的，因此效率很低。MobileDiffusion 提出直接將 softmax 函數 finetune 到 relu，因為 relu 是每一個點的激活，更為高效。令人驚訝的是，只需要大概一萬步的微調，模型 metric 反而提升了，出圖質量也有保障。因此 relu 相比於 softmax 的優勢是很明顯的了。
Separable Convolution ( 可分離卷積 )
MobileDiffuison 精簡參數的關鍵還在採用了 Seprable Convolution。這種技術已經被 MobileNet 等工作證實是極為有效的，特別是移動端，但是一般在生成模型上很少採用。MobileDiffusion 實驗發現 Separable Convolution 對減少參數是很有效的，尤其是將其放在 UNet 最內層，模型質量經分析證明是沒有損失的。
採樣優化
當下最常採用的採樣優化方法包括了 Progressive Distillation 和 UFOGen， 分別可以做到 8 steps 和 1 step。為了證明在模型極致精簡後，這些採樣依然適用，MobileDiffusion 對兩者同時做了實驗驗證。
採樣優化前後和基準模型的比較如下，可以看出採樣優化後的 8 steps 和 1 step 的模型，指標都是比較突出的。
實驗與應用
移動端基準測試
MobileDiffusion 在 iPhone 15 Pro 上可以得到當前最快的出圖速度，0.2s！
下游任務測試
MobileDiffusion 探索了包括 ControlNet/Plugin 和 LoRA Finetune 的下游任務。從下圖可以看出，經過模型和採樣優化後，MobileDiffusion 依然保持了優秀的模型微調能力。
總結
MobileDiffusion 探索了多種模型和採樣優化方法，最終可以實現在移動端的亞秒級出圖能力，下游微調應用依然有保障。我們相信這將會對今後高效的擴散模型設計產生影響，並拓展移動端應用實例。
-
-
-
-
-
-
-
-
-
-