手機(jī)上0.2秒出圖、當(dāng)前速度之最，谷歌打造超快擴(kuò)散模型MobileDiffusion

　　本文來自于微信公眾號(hào) 機(jī)器之心(ID:almosthuman2014)，作者:機(jī)器之心。

　　在手機(jī)等移動(dòng)端側(cè)運(yùn)行 Stable Diffusion 等文生圖生成式 AI 大模型已經(jīng)成為業(yè)界追逐的熱點(diǎn)之一，其中生成速度是主要的制約因素。

　　近日，來自谷歌的一篇論文「MobileDiffusion: Subsecond Text-to-Image Generation on Mobile Devices」，提出了手機(jī)端最快文生圖，在 iPhone15Pro 上只要0.2秒。論文出自 UFOGen 同一團(tuán)隊(duì)，在打造超小擴(kuò)散模型的同時(shí)， 采用當(dāng)前大火的 Diffusion GAN 技術(shù)路線做采樣加速。

　　下面是 MobileDiffusion 一步生成的結(jié)果。

　　那么，MobileDiffusion 是如何優(yōu)化得到的呢?

　　讓我們先從問題出發(fā)，為什么優(yōu)化是必須的。

　　當(dāng)下最火的文本到圖像生成都是基于擴(kuò)散模型來實(shí)現(xiàn)的。依賴于其預(yù)訓(xùn)練模型強(qiáng)大的基本圖像生成能力和在下游微調(diào)任務(wù)上的穩(wěn)健性質(zhì)， 我們看到了擴(kuò)散模型在諸如圖像編輯、可控生成、 個(gè)性化生成以及視頻生成的非凡表現(xiàn)。

　　然而作為 Foundation Model， 它的不足也很明顯，主要包括了兩方面:一是擴(kuò)散模型的大量參數(shù)導(dǎo)致計(jì)算速度慢，尤其是在資源有限的情況下;二是擴(kuò)散模型需要多步才能采樣，這進(jìn)一步導(dǎo)致很慢的推理速度。拿最火的的 Stable Diffusion1.5(SD) 為例，其基礎(chǔ)模型包含了接近10億參數(shù)，我們?cè)?iPhone15Pro 上將模型量化后進(jìn)行推理，50步采樣需要接近80s。如此昂貴的資源需求和遲滯的用戶體驗(yàn)極大的限制了其在移動(dòng)端的應(yīng)用場景。

　　為了解決以上問題，MobileDiffusion 點(diǎn)對(duì)點(diǎn)地進(jìn)行優(yōu)化。(1) 針對(duì)模型體積龐大的問題，我們主要對(duì)其核心組件 UNet 進(jìn)行了大量試驗(yàn)及優(yōu)化，包括了將計(jì)算昂貴的卷積精簡和注意力運(yùn)算放在了較低的層上，以及針對(duì) Mobile Devices 的操作優(yōu)化，諸如激活函數(shù)等。(2)針對(duì)擴(kuò)散模型需要多步采樣的問題， MobileDiffusion 探索并實(shí)踐了像 Progressive Distillation 和當(dāng)前最先進(jìn)的 UFOGen 的一步推理技術(shù)。

　　模型優(yōu)化

　　MobileDiffusion 基于當(dāng)下開源社區(qū)里最火的 SD1.5UNet 進(jìn)行優(yōu)化。在每次的優(yōu)化操作后， 會(huì)同時(shí)衡量相對(duì)于原始 UNet 模型的性能的損失，測量指標(biāo)包括 FID 和 CLIP 兩個(gè)常用 metric。

　　宏觀設(shè)計(jì)

　　上圖左邊是原始 UNet 的設(shè)計(jì)示意， 可以看出基本包括了 Convolution 和 Transformer，Transformer 又包括了 Self-Attention 和 Cross-Attention。

　　MobileDiffusion 對(duì) UNet 優(yōu)化的核心思路分為兩點(diǎn):1)精簡 Convolution， 眾所周知，在高分辨率的特征空間上進(jìn)行了 Convolution 是十分耗時(shí)的， 而且參數(shù)量很大，這里指的是 Full Convolution;2)提高 Attention 效率。和 Convolution 一樣，高 Attention 需要對(duì)整個(gè)特征空間的長度進(jìn)行運(yùn)算，Self-Attention 復(fù)雜度和特征空間展平后長度成平方關(guān)系，Cross-Attention 也要和空間長度成正比。

　　實(shí)驗(yàn)表明將整個(gè) UNet 的16個(gè) Transformer 移到特征分辨率最低的內(nèi)層，同時(shí)每一層都剪掉一個(gè)卷積，不會(huì)對(duì)性能有明顯影響。達(dá)到的效果就是:MobileDiffusion 將原本22個(gè) Convolution 和16個(gè) Transformer，可以極致地精簡到11個(gè) Convolution 和12個(gè)左右 Transformer，并且這些注意力都是在低分辨率特征圖上進(jìn)行的，因?yàn)樾�率�?huì)極大提升，帶來了40% 效率提升和40% 參數(shù)剪切，最終模型如上圖右圖所示。和更多模型的對(duì)比如下:

　　微觀設(shè)計(jì)

　　這里將只介紹幾種新穎的設(shè)計(jì)，有興趣的讀者可以閱讀正文， 會(huì)有更詳細(xì)的介紹。

　　Decouple Self-Attention and Cross-Attention

　　傳統(tǒng) UNet 里 Transformer 同時(shí)包含 Self-Attention 和 Cross-Attention，MobileDiffusion 將 Self-Attention 全部放在了最低分辨率特征圖，但是保留一個(gè) Cross-Attention 在中間層，發(fā)現(xiàn)這種設(shè)計(jì)既提高了運(yùn)算效率又保證了模型出圖質(zhì)量

　　Finetune softmax into relu

　　softmax 眾所周知在大部分未優(yōu)化情況下是難以并行的，因此效率很低。MobileDiffusion 提出直接將 softmax 函數(shù) finetune 到 relu，因?yàn)?relu 是每一個(gè)點(diǎn)的激活，更為高效。令人驚訝的是，只需要大概一萬步的微調(diào)，模型 metric 反而提升了，出圖質(zhì)量也有保障。因此 relu 相比于 softmax 的優(yōu)勢是很明顯的了。

　　Separable Convolution (可分離卷積)

　　MobileDiffuison 精簡參數(shù)的關(guān)鍵還在采用了 Seprable Convolution。這種技術(shù)已經(jīng)被 MobileNet 等工作證實(shí)是極為有效的，特別是移動(dòng)端，但是一般在生成模型上很少采用。MobileDiffusion 實(shí)驗(yàn)發(fā)現(xiàn) Separable Convolution 對(duì)減少參數(shù)是很有效的，尤其是將其放在 UNet 最內(nèi)層，模型質(zhì)量經(jīng)分析證明是沒有損失的。

　　采樣優(yōu)化

　　當(dāng)下最常采用的采樣優(yōu)化方法包括了 Progressive Distillation 和 UFOGen， 分別可以做到8steps 和1step。為了證明在模型極致精簡后，這些采樣依然適用，MobileDiffusion 對(duì)兩者同時(shí)做了實(shí)驗(yàn)驗(yàn)證。

　　采樣優(yōu)化前后和基準(zhǔn)模型的比較如下，可以看出采樣優(yōu)化后的8steps 和1step 的模型，指標(biāo)都是比較突出的。

　　實(shí)驗(yàn)與應(yīng)用

　　移動(dòng)端基準(zhǔn)測試

　　MobileDiffusion 在 iPhone15Pro 上可以得到當(dāng)前最快的出圖速度，0.2s!

　　下游任務(wù)測試

　　MobileDiffusion 探索了包括 ControlNet/Plugin 和 LoRA Finetune 的下游任務(wù)。從下圖可以看出，經(jīng)過模型和采樣優(yōu)化后，MobileDiffusion 依然保持了優(yōu)秀的模型微調(diào)能力。

　　總結(jié)

　　MobileDiffusion 探索了多種模型和采樣優(yōu)化方法，最終可以實(shí)現(xiàn)在移動(dòng)端的亞秒級(jí)出圖能力，下游微調(diào)應(yīng)用依然有保障。我們相信這將會(huì)對(duì)今后高效的擴(kuò)散模型設(shè)計(jì)產(chǎn)生影響，并拓展移動(dòng)端應(yīng)用實(shí)例。