Stable Diffusion 在圖像生成領(lǐng)域的知名度不亞于對話大模型中的 ChatGPT。其能夠在幾十秒內(nèi)為任何給定的輸入文本創(chuàng)建逼真圖像。由于 Stable Diffusion 的參數(shù)量超過 10 億，并且由于設(shè)備上的計(jì)算和內(nèi)存資源有限，因而這種模型主要運(yùn)行在云端。

在沒有精心設(shè)計(jì)和實(shí)施的情況下，在設(shè)備上運(yùn)行這些模型可能會(huì)導(dǎo)致延遲增加，這是由于迭代降噪過程和內(nèi)存消耗過多造成的。

如何在設(shè)備端運(yùn)行 Stable Diffusion 引起了大家的研究興趣，此前，有研究者開發(fā)了一個(gè)應(yīng)用程序，該應(yīng)用在 iPhone 14 Pro 上使用 Stable Diffusion 生成圖片僅需一分鐘，使用大約 2GiB 的應(yīng)用內(nèi)存。

此前蘋果也對此做了一些優(yōu)化，他們在 iPhone、iPad、Mac 等設(shè)備上，半分鐘就能生成一張分辨率 512x512 的圖像。高通緊隨其后，在安卓手機(jī)端運(yùn)行 Stable Diffusion v1.5 ，不到 15 秒生成分辨率 512x512 的圖像。

近日，谷歌發(fā)表的一篇論文中《 Speed Is All You Need: On-Device Acceleration of Large Diffusion Models via GPU-Aware Optimizations 》，他們實(shí)現(xiàn)了在 GPU 驅(qū)動(dòng)的設(shè)備上運(yùn)行 Stable Diffusion 1.4 ，達(dá)到 SOTA 推理延遲性能（在三星 S23 Ultra 上，通過 20 次迭代生成 512 × 512 的圖像僅需 11.5 秒）。此外，該研究不是只針對一種設(shè)備；相反，它是一種通用方法，適用于改進(jìn)所有潛在擴(kuò)散模型。

在沒有數(shù)據(jù)連接或云服務(wù)器的情況下，這項(xiàng)研究為在手機(jī)上本地運(yùn)行生成 AI 開辟了許多可能性。Stable Diffusion 去年秋天才發(fā)布，今天已經(jīng)可以塞進(jìn)設(shè)備運(yùn)行，可見這個(gè)領(lǐng)域發(fā)展速度有多快。

論文地址：https://arxiv.org/pdf/2304.11267.pdf

為了達(dá)到這一生成速度，谷歌提出了一些優(yōu)化建議，下面我們看看谷歌是如何優(yōu)化的。

方法介紹

該研究旨在提出優(yōu)化方法來提高大型擴(kuò)散模型文生圖的速度，其中針對 Stable Diffusion 提出一些優(yōu)化建議，這些優(yōu)化建議也適用于其他大型擴(kuò)散模型。

首先來看一下 Stable Diffusion 的主要組成部分，包括：文本嵌入器（text embedder）、噪聲生成（noise generation）、去噪神經(jīng)網(wǎng)絡(luò)（denoising neural network）和圖像解碼器（image decoder，如下圖 1 所示。

然后我們具體看一下該研究提出的三種優(yōu)化方法。

專用內(nèi)核：Group Norm 和 GELU

組歸一化（GN）方法的工作原理是將特征圖的通道（channel）劃分為更小的組，并獨(dú)立地對每個(gè)組進(jìn)行歸一化，從而使 GN 對批大小的依賴性降低，更適合各種批大小和網(wǎng)絡(luò)架構(gòu)。該研究沒有按順序執(zhí)行 reshape、取均值、求方差、歸一化這些操作，而是設(shè)計(jì)了一個(gè)獨(dú)特的 GPU shader 形式的內(nèi)核，它可以在一個(gè) GPU 命令中執(zhí)行所有這些操作，而無需任何中間張量（tensor）。

高斯誤差線性單元（GELU）作為常用的模型激活函數(shù)，包含大量數(shù)值計(jì)算，例如乘法、加法和高斯誤差函數(shù)。該研究用一個(gè)專用的 shader 來整合這些數(shù)值計(jì)算及其伴隨的 split 和乘法操作，使它們能夠在單個(gè) AI 作畫調(diào)用中執(zhí)行。

提高注意力模塊的效率

Stable Diffusion 中的文本到圖像 transformer 有助于對條件分布進(jìn)行建模，這對于文本到圖像生成任務(wù)至關(guān)重要。然而，由于內(nèi)存復(fù)雜性和時(shí)間復(fù)雜度，自 / 交叉注意力機(jī)制在處理長序列時(shí)遇到了困難?；诖?，該研究提出兩種優(yōu)化方法，以緩解計(jì)算瓶頸。

一方面，為了避免在大矩陣上執(zhí)行整個(gè) softmax 計(jì)算，該研究使用一個(gè) GPU shader 來減少運(yùn)算操作，大大減少了中間張量的內(nèi)存占用和整體延遲，具體方法如下圖 2 所示。

另一方面，該研究采用 FlashAttention [7] 這種 IO 感知的精確注意力算法，使得高帶寬內(nèi)存（HBM）的訪問次數(shù)少于標(biāo)準(zhǔn)注意力機(jī)制，提高了整體效率。

Winograd 卷積

Winograd 卷積將卷積運(yùn)算轉(zhuǎn)換為一系列矩陣乘法。這種方法可以減少許多乘法運(yùn)算，提高計(jì)算效率。但是，這樣一來也會(huì)增加內(nèi)存消耗和數(shù)字錯(cuò)誤，特別是在使用較大的 tile 時(shí)。

Stable Diffusion 的主干在很大程度上依賴于 3×3 卷積層，尤其是在圖像解碼器中，它們占了 90% 。該研究對這一現(xiàn)象進(jìn)行了深入分析，以探索在 3 × 3 內(nèi)核卷積上使用不同 tile 大小的 Winograd 的潛在好處。研究發(fā)現(xiàn) 4 × 4 的 tile 大小最佳，因?yàn)樗谟?jì)算效率和內(nèi)存利用率之間提供了最佳平衡。

實(shí)驗(yàn)

該研究在各種設(shè)備上進(jìn)行了基準(zhǔn)測試：三星 S23 Ultra（Adreno 740）和 iPhone 14 Pro Max（A16）。基準(zhǔn)測試結(jié)果如下表 1 所示：

很明顯，隨著每個(gè)優(yōu)化被激活，延遲逐漸減少（可理解為生成圖像時(shí)間減少）。具體而言，與基線相比：在三星 S23 Ultra 延遲減少 52.2%；iPhone 14 Pro Max 延遲減少 32.9%。此外，該研究還對三星 S23 Ultra 端到端延遲進(jìn)行評估，在 20 個(gè)去噪迭代 step 內(nèi)，生成 512 × 512 像素圖像，不到 12 秒就達(dá)到 SOTA 結(jié)果。

小型設(shè)備可以運(yùn)行自己的生成式人工智能模型，這對未來意味著什么？我們可以期待一波。