原文來(lái)源：機(jī)器之心

圖片來(lái)源：由無(wú)界 AI生成

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

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

論文地址：https://arxiv.org/abs/2311.16567

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

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

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

當(dāng)下最火的文本到圖像生成都是基于擴(kuò)散模型來(lái)實(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 Diffusion 1.5 (SD) 為例,，其基礎(chǔ)模型包含了接近 10 億參數(shù)，我們?cè)?iPhone 15 Pro 上將模型量化后進(jìn)行推理,，50 步采樣需要接近 80s,。如此昂貴的資源需求和遲滯的用戶體驗(yàn)極大的限制了其在移動(dòng)端的應(yīng)用場(chǎng)景,。

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

模型優(yōu)化

MobileDiffusion 基于當(dāng)下開源社區(qū)里最火的 SD 1.5 UNet 進(jìn)行優(yōu)化。在每次的優(yōu)化操作后,， 會(huì)同時(shí)衡量相對(duì)于原始 UNet 模型的性能的損失,，測(cè)量指標(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）精簡(jiǎn) Convolution， 眾所周知,，在高分辨率的特征空間上進(jìn)行了 Convolution 是十分耗時(shí)的,， 而且參數(shù)量很大，這里指的是 Full Convolution,；2）提高 Attention 效率,。和 Convolution 一樣，高 Attention 需要對(duì)整個(gè)特征空間的長(zhǎng)度進(jìn)行運(yùn)算,，Self-Attention 復(fù)雜度和特征空間展平后長(zhǎng)度成平方關(guān)系,，Cross-Attention 也要和空間長(zhǎng)度成正比。

實(shí)驗(yàn)表明將整個(gè) UNet 的 16 個(gè) Transformer 移到特征分辨率最低的內(nèi)層,，同時(shí)每一層都剪掉一個(gè)卷積,，不會(huì)對(duì)性能有明顯影響。達(dá)到的效果就是：MobileDiffusion 將原本 22 個(gè) Convolution 和 16 個(gè) Transformer,，可以極致地精簡(jiǎn)到 11 個(gè) Convolution 和 12 個(gè)左右 Transformer,，并且這些注意力都是在低分辨率特征圖上進(jìn)行的，因?yàn)樾蕰?huì)極大提升,，帶來(lái)了 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)的激活,，更為高效。令人驚訝的是,，只需要大概一萬(wàn)步的微調(diào),，模型 metric 反而提升了,，出圖質(zhì)量也有保障,。因此 relu 相比于 softmax 的優(yōu)勢(shì)是很明顯的了。

Separable Convolution (可分離卷積)

MobileDiffuison 精簡(jiǎn)參數(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)分析證明是沒(méi)有損失的,。

采樣優(yōu)化

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

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

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

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

MobileDiffusion 在 iPhone 15 Pro 上可以得到當(dāng)前最快的出圖速度，0.2s,！

下游任務(wù)測(cè)試

MobileDiffusion 探索了包括 ControlNet/Plugin 和 LoRA Finetune 的下游任務(wù),。從下圖可以看出,，經(jīng)過(guò)模型和采樣優(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í)例,。