編輯：桃子 拉燕

【新智元導讀】手機12秒一鍵生圖？谷歌最新研究做到了。

手機生圖只要12秒？

這不是吹的，谷歌就給實現了。

最新研究中，谷歌研究人員對4個GPU分層優化，在三星手機上成功運行Stable Diffusion 1.4。

實現了11.5秒生成圖像，重要的是，內存使用量也大幅減少。

正所謂，Speed Is All You Need！

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

谷歌最新提出的方法是通用的，可以對所有擴散模型的改進，並非僅針對特定設備實現的。

通過實驗結果發現，三星S23 Ultra和iPhone 14 Pro的整體圖像生成時間分別減少了52%和33%。

這意味著，一部手機裝下一個生成式AI模型的未來更近了。

從3080到一部手機

當前，將大型擴散模型合併到任何APP中的一個關鍵考慮因素是，模型將在何處執行選擇。

在一個消費級設備上部署模型的好處是，較低的服務成本、改善擴展性、還可以離線，並且能改善用戶隱私。

22年，Stable Diffusion剛剛發布的第一個版本，最初只能緩慢運行在RTX 3080上。

Stable Diffusion有超過10億的參數，DALL-E是120億，以後隨著擴散模型的發展，參數量會逐漸增加。

由於設備計算和內存資源的限制，因此在運行時帶來了諸多挑戰。

在沒有精心設計情況下，在設備上運行這些模型可能會導致，由於疊代去噪過程和過多的內存消耗，輸出內容延遲就會增加。

此前，也有研究實現了將Stable Diffusion成功部署到設備上，但僅局限於特定的設備或晶片組。

對此，谷歌研究人員對大型擴散模型提供了一系列實現優化，這些模型在配備 GPU 的行動裝置上實現了迄今為止報導的最快推理延遲。

在不使用INT8量化的情況下，對於一張512x512的圖片進行20次疊代，Stable Diffusion 1.4的推理延遲低於12秒。

具體是如何實現的呢？

GPU感知優化

在論文中，研究人員側重的是使用大型擴散模型，然後完成從文本描述生成圖像的任務。

雖說論文中，部分討論是研究者為Stable Diffusion特定結構所提出的優化建議，但這些優化可以很容易推廣到其它大型擴散模型上。

研究人員表示，當用文本提示進行推理時，這個過程包含根據所需的文本描述，應用額外條件來指導反向擴散。

具體來說，Stable Diffusion的主要組成部分包括：文本嵌入器（Text Embedder）、噪聲生成（Noise Generation）、去噪神經網絡（Denoising Neural Network，aka UNet），以及圖像解碼器（Image Decoder）。

如下圖所示：

Stable Diffusion中主要組件及其相互作用的示意圖

下面，我們分別介紹一下這幾個組成部分，各部分間的關係參照圖。

· 文本嵌入器：

利用CLIP模型對文本提示y進行編碼，生成一個高維嵌入向量τθ(y)，將文本提示的語義封裝進去。該嵌入被當作去噪神經網絡的輸入，為逆向擴散的過程提供指示。

· 噪聲生成：

給潛在空間提供隨機噪聲z，該噪聲作為逆向擴散過程的起始點。

· 去噪神經網絡：

該網絡被設計為近似p(z|y)形式的條件分布，利用條件去噪自動編碼器θ(zt, t, τθ(y))（denoising autoencoder）。每次疊代t採用UNet架構。

同時，交叉注意機制（cross-attention mechanism）被用來操作潛在空間和文本嵌入向量，在疊代過程中預測z的去噪版本。

· 圖像解碼器：

• 逆行擴散過程在潛在空間
• 中進行。一旦這個過程完成，圖像解碼器D被用來從潛在矢量中重建RGB圖像。
• 
  

研究人員在整個UNet架構中實現了群組歸一化（Group normalization，GN）。

這種歸一化技術的工作原理是將特徵圖（feature map）的pipeline劃分為較小的組，並對每個組進行獨立的歸一化，使GN對批次大小的依賴性降低，更適合於各種大小的批次和各種網絡結構。

應用公式①，每個特徵值

被歸一化為其所屬組的組均值

和方差

。

（公式①）

研究人員並沒有依次執行上述提到的重塑、平均值、方差和歸一化的所有操作，而是以GPU Shader的形式設計了一個特別的程序，在一個GPU命令中執行所有這些操作，無需中間流程。

這裡先介紹一下Gaussian Error Linear Unit（GELU）。

GELU作為模型中普遍存在的激活函數，包含許多數值計算，如乘法、加法和高斯誤差函數，如公式②所示。

研究人員弄了一個專門的Shader來整合這些數值計算及其伴隨的分割和乘法操作，使其在一次繪圖調用中完成執行。

（公式②）

穩定擴散中的文本/圖像變換器有助於對條件分布P(z|τθ(y))進行建模，這對文本到圖像的生成任務至關重要。

然而，自我/交叉注意力機制在處理長序列時遇到了困難，因為它們的時間和內存複雜性是平過方的。在論文中，研究人員介紹了兩種可能的優化，旨在緩解這些計算瓶頸。

一種是Partially Fused Softmax，另一種是FlashAttention。

下面僅以Softmax為例。

上圖是在注意力模塊中，優化過的softmax實現。

虛線以上的流程圖描述的是直接在矩陣

中應用softmax的初始實現。

虛線以下的展示的則是修正後的模塊（紅色部分）。

總而言之，論文中研究人員提出了一整套優化方案，可以在各種設備上執行大型擴散模型時，共同達到了突破性的延遲數字。

這些改進擴大了模型的通用性，並提高了在各種設備上的整體用戶體驗。

12秒，業界領先

為了評估改進後的模型，研究人員分別在三星S23 Ultra (Adreno 740) 和iPhone 14 Pro Max (A16) 進行了一組基準測試。

作為去噪神經網絡，UNet是計算需求最高的組件。

研究人員提供了執行單次疊代的UNet所需的延遲數據，以毫秒為單位測量，圖像解析度為512x512。

此外，他們記錄了運行時生成的中間張量在「Tensor」列中的內存使用情況，以及為保存模型權重分配的內存在「Weight」列中的使用情況，均以兆字節為單位。

請注意，內存管理器通過重用中間張量的緩衝區來優化內存占用。

如表中數據顯示，第一行顯示了在公共Github倉庫中使用內部OpenCL內核實現，但沒有任何優化的結果。

實現之後的結果，並且研究者在沒有任何優化的情況下使用內部OpenCL內核。

第2-5行，分別逐個啟用每個優化：

Opt. Softmax：部分融合的softmax和優化的softmax減少步驟

S-GN/GELU：用於組歸一化和GELU的專用內核

FlashAttn.：FlashAttention實現

Winograd(All)：採用Winograd卷積

隨著每個優化的啟用，實驗結果發現延遲逐步減少。

與基線相比，在兩種設備上都觀察到了顯著的總體延遲降低：三星S23 Ultra 降低52.2%，iPhone 14 Pro Max降低32.9%。

此外，研究人員還評估了在三星S23 Ultra進行文本到圖像輸出端到端延遲。

進行了20次去噪疊代，生成一張512x512圖像，實現了不到12秒的業界領先結果。

可見，在沒有數據連接或雲伺服器的情況下，在手機上本地運行生成式人工智慧模型，將開闢了許多可能性。

谷歌最新研究給出了一種全新方案。

參考資料：

https://arxiv.org/abs/2304.11267

https://www.reddit.com/r/MachineLearning/comments/12zclus/d_google_researchers_achieve_performance/