谷歌下場優(yōu)化擴散模型，三星手機運行 Stable Diffusion，12 秒內(nèi)出圖

機器之心報道

編輯：陳萍、小舟

(資料圖片僅供參考)

Speed Is All You Need：谷歌提出針對 Stable Diffusion 一些優(yōu)化建議，生成圖片速度快速提升。

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

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

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

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

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

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

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

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

方法介紹

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

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

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

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

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

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

提高注意力模塊的效率

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

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

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

Winograd 卷積

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

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

實驗

該研究在各種設備上進行了基準測試：三星 S23 Ultra（Adreno 740）和 iPhone 14 Pro Max（A16）?；鶞蕼y試結(jié)果如下表 1 所示：

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

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