速度驚人，手機跑 Stable Diffusion，12 秒出圖，谷歌加速擴散模型破記錄

手機生圖只要 12 秒？

這不是吹的，谷歌就給實現(xiàn)了。

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

實現(xiàn)了 11.5 秒生成圖像，重要的是，內(nèi)存使用量也大幅減少。

正所謂，Speed Is All You Need！

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

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

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

從 3080 到一部手機

當前，將大型擴散模型合并到任何 App 中的一個關鍵考慮因素是，模型將在何處執(zhí)行選擇。

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

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

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

由于設備計算和內(nèi)存資源的限制，因此在運行時帶來了諸多挑戰(zhàn)。

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

此前，也有研究實現(xiàn)了將 Stable Diffusion 成功部署到設備上，但僅局限于特定的設備或芯片組。

對此，谷歌研究人員對大型擴散模型提供了一系列實現(xiàn)優(yōu)化，這些模型在配備 GPU 的移動設備上實現(xiàn)了迄今為止報道的最快推理延遲。

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

具體是如何實現(xiàn)的呢？

GPU 感知優(yōu)化

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

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

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

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

如下圖所示：

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

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

?文本嵌入器：

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

?噪聲生成：

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

?去噪神經(jīng)網(wǎng)絡：

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

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

?圖像解碼器：

逆行擴散過程在潛在空間

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

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

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

應用公式①，每個特征值

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

和方差

。

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

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

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

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

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

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

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

下面僅以 Softmax 為例。

上圖是在注意力模塊中，優(yōu)化過的 softmax 實現(xiàn)。

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

中應用 softmax 的初始實現(xiàn)。

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

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

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

12 秒，業(yè)界領先

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

作為去噪神經(jīng)網(wǎng)絡，UNet 是計算需求最高的組件。

研究人員提供了執(zhí)行單次迭代的 UNet 所需的延遲數(shù)據(jù)，以毫秒為單位測量，圖像分辨率為 512x512。

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

請注意，內(nèi)存管理器通過重用中間張量的緩沖區(qū)來優(yōu)化內(nèi)存占用。

如表中數(shù)據(jù)顯示，第一行顯示了在公共 Github 倉庫中使用內(nèi)部 OpenCL 內(nèi)核實現(xiàn)，但沒有任何優(yōu)化的結果。

實現(xiàn)之后的結果，并且研究者在沒有任何優(yōu)化的情況下使用內(nèi)部 OpenCL 內(nèi)核。

第 2-5 行，分別逐個啟用每個優(yōu)化：

Opt. Softmax：部分融合的 softmax 和優(yōu)化的 softmax 減少步驟

S-GN / GELU：用于組歸一化和 GELU 的專用內(nèi)核

FlashAttn.：FlashAttention 實現(xiàn)

Winograd (All)：采用 Winograd 卷積

隨著每個優(yōu)化的啟用，實驗結果發(fā)現(xiàn)延遲逐步減少。

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

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

進行了 20 次去噪迭代，生成一張 512x512 圖像，實現(xiàn)了不到 12 秒的業(yè)界領先結果。

可見，在沒有數(shù)據(jù)連接或云服務器的情況下，在手機上本地運行生成式人工智能模型，將開辟了許多可能性。

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

參考資料：

• https://arxiv.org/abs/2304.11267

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

本文來自微信公眾號：新智元 （ID：AI_era）

廣告聲明：文內(nèi)含有的對外跳轉鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時間，結果僅供參考，IT之家所有文章均包含本聲明。