手機 12 秒一鍵生圖？谷歌最新研究做到了。
手機生圖只要 12 秒？
這不是吹的，谷歌就給實現(xiàn)了。
最新研究中，谷歌研究人員對 4 個 GPU 分層優(yōu)化，在三星手機上成功運行 Stable Diffusion 1.4。
實現(xiàn)了 11.5 秒生成圖像，重要的是，內(nèi)存使用量也大幅減少。
正所謂，Speed Is All You Need！
谷歌最新提出的方法是通用的，可以對所有擴散模型的改進，并非僅針對特定設(shè)備實現(xiàn)的。
通過實驗結(jié)果發(fā)現(xiàn)，三星 S23 Ultra 和 iPhone 14 Pro 的整體圖像生成時間分別減少了 52% 和 33%。
這意味著，一部手機裝下一個生成式 AI 模型的未來更近了。
從 3080 到一部手機
當(dāng)前，將大型擴散模型合并到任何 App 中的一個關(guān)鍵考慮因素是，模型將在何處執(zhí)行選擇。
在一個消費級設(shè)備上部署模型的好處是，較低的服務(wù)成本、改善擴展性、還可以離線，并且能改善用戶隱私。
22 年，Stable Diffusion 剛剛發(fā)布的第一個版本，最初只能緩慢運行在 RTX 3080 上。
Stable Diffusion 有超過 10 億的參數(shù)，DALL-E 是 120 億，以后隨著擴散模型的發(fā)展，參數(shù)量會逐漸增加。
由于設(shè)備計算和內(nèi)存資源的限制，因此在運行時帶來了諸多挑戰(zhàn)。
在沒有精心設(shè)計情況下，在設(shè)備上運行這些模型可能會導(dǎo)致，由于迭代去噪過程和過多的內(nèi)存消耗，輸出內(nèi)容延遲就會增加。
此前，也有研究實現(xiàn)了將 Stable Diffusion 成功部署到設(shè)備上，但僅局限于特定的設(shè)備或芯片組。
對此，谷歌研究人員對大型擴散模型提供了一系列實現(xiàn)優(yōu)化，這些模型在配備 GPU 的移動設(shè)備上實現(xiàn)了迄今為止報道的最快推理延遲。
在不使用 INT8 量化的情況下，對于一張 512x512 的圖片進行 20 次迭代，Stable Diffusion 1.4 的推理延遲低于 12 秒。
具體是如何實現(xiàn)的呢？
GPU 感知優(yōu)化
在論文中，研究人員側(cè)重的是使用大型擴散模型，然后完成從文本描述生成圖像的任務(wù)。
雖說論文中，部分討論是研究者為 Stable Diffusion 特定結(jié)構(gòu)所提出的優(yōu)化建議，但這些優(yōu)化可以很容易推廣到其它大型擴散模型上。
研究人員表示，當(dāng)用文本提示進行推理時，這個過程包含根據(jù)所需的文本描述，應(yīng)用額外條件來指導(dǎo)反向擴散。
具體來說，Stable Diffusion 的主要組成部分包括：文本嵌入器（Text Embedder）、噪聲生成（Noise Generation）、去噪神經(jīng)網(wǎng)絡(luò)（Denoising Neural Network，aka UNet），以及圖像解碼器（Image Decoder）。
如下圖所示：
Stable Diffusion 中主要組件及其相互作用的示意圖
下面，我們分別介紹一下這幾個組成部分，各部分間的關(guān)系參照圖。
?文本嵌入器：
利用 CLIP 模型對文本提示 y 進行編碼，生成一個高維嵌入向量 τθ(y)，將文本提示的語義封裝進去。該嵌入被當(dāng)作去噪神經(jīng)網(wǎng)絡(luò)的輸入，為逆向擴散的過程提供指示。
?噪聲生成：
給潛在空間提供隨機噪聲 z，該噪聲作為逆向擴散過程的起始點。
?去噪神經(jīng)網(wǎng)絡(luò)：
該網(wǎng)絡(luò)被設(shè)計為近似 p (z|y) 形式的條件分布，利用條件去噪自動編碼器 θ(zt, t, τθ(y))（denoising autoencoder）。每次迭代 t 采用 UNet 架構(gòu)。
同時，交叉注意機制（cross-attention mechanism）被用來操作潛在空間和文本嵌入向量，在迭代過程中預(yù)測 z 的去噪版本。
?圖像解碼器：
逆行擴散過程在潛在空間
中進行。一旦這個過程完成，圖像解碼器 D 被用來從潛在矢量中重建 RGB 圖像。
研究人員在整個 UNet 架構(gòu)中實現(xiàn)了群組歸一化（Group normalization，GN）。
這種歸一化技術(shù)的工作原理是將特征圖（feature map）的 pipeline 劃分為較小的組，并對每個組進行獨立的歸一化，使 GN 對批次大小的依賴性降低，更適合于各種大小的批次和各種網(wǎng)絡(luò)結(jié)構(gòu)。
應(yīng)用公式①，每個特征值
被歸一化為其所屬組的組均值
和方差
。
研究人員并沒有依次執(zhí)行上述提到的重塑、平均值、方差和歸一化的所有操作，而是以 GPU Shader 的形式設(shè)計了一個特別的程序，在一個 GPU 命令中執(zhí)行所有這些操作，無需中間流程。
這里先介紹一下 Gaussian Error Linear Unit（GELU）。
GELU 作為模型中普遍存在的激活函數(shù)，包含許多數(shù)值計算，如乘法、加法和高斯誤差函數(shù)，如公式②所示。
研究人員弄了一個專門的 Shader 來整合這些數(shù)值計算及其伴隨的分割和乘法操作，使其在一次繪圖調(diào)用中完成執(zhí)行。
穩(wěn)定擴散中的文本 / 圖像變換器有助于對條件分布 P (z|τθ(y)) 進行建模，這對文本到圖像的生成任務(wù)至關(guān)重要。
然而，自我 / 交叉注意力機制在處理長序列時遇到了困難，因為它們的時間和內(nèi)存復(fù)雜性是平過方的。在論文中，研究人員介紹了兩種可能的優(yōu)化，旨在緩解這些計算瓶頸。
一種是 Partially Fused Softmax，另一種是 FlashAttention。
下面僅以 Softmax 為例。
上圖是在注意力模塊中，優(yōu)化過的 softmax 實現(xiàn)。
虛線以上的流程圖描述的是直接在矩陣
中應(yīng)用 softmax 的初始實現(xiàn)。
虛線以下的展示的則是修正后的模塊（紅色部分）。
總而言之，論文中研究人員提出了一整套優(yōu)化方案，可以在各種設(shè)備上執(zhí)行大型擴散模型時，共同達到了突破性的延遲數(shù)字。
這些改進擴大了模型的通用性，并提高了在各種設(shè)備上的整體用戶體驗。
12 秒，業(yè)界領(lǐng)先
為了評估改進后的模型，研究人員分別在三星 S23 Ultra (Adreno 740) 和 iPhone 14 Pro Max (A16) 進行了一組基準(zhǔn)測試。
作為去噪神經(jīng)網(wǎng)絡(luò)，UNet 是計算需求最高的組件。
研究人員提供了執(zhí)行單次迭代的 UNet 所需的延遲數(shù)據(jù)，以毫秒為單位測量，圖像分辨率為 512x512。
此外，他們記錄了運行時生成的中間張量在「Tensor」列中的內(nèi)存使用情況，以及為保存模型權(quán)重分配的內(nèi)存在「Weight」列中的使用情況，均以兆字節(jié)為單位。
請注意，內(nèi)存管理器通過重用中間張量的緩沖區(qū)來優(yōu)化內(nèi)存占用。
如表中數(shù)據(jù)顯示，第一行顯示了在公共 Github 倉庫中使用內(nèi)部 OpenCL 內(nèi)核實現(xiàn)，但沒有任何優(yōu)化的結(jié)果。
實現(xiàn)之后的結(jié)果，并且研究者在沒有任何優(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)化的啟用，實驗結(jié)果發(fā)現(xiàn)延遲逐步減少。
與基線相比，在兩種設(shè)備上都觀察到了顯著的總體延遲降低：三星 S23 Ultra 降低 52.2%，iPhone 14 Pro Max 降低 32.9%。
此外，研究人員還評估了在三星 S23 Ultra 進行文本到圖像輸出端到端延遲。
進行了 20 次去噪迭代，生成一張 512x512 圖像，實現(xiàn)了不到 12 秒的業(yè)界領(lǐng)先結(jié)果。
可見，在沒有數(shù)據(jù)連接或云服務(wù)器的情況下，在手機上本地運行生成式人工智能模型，將開辟了許多可能性。
谷歌最新研究給出了一種全新方案。
參考資料：
https://arxiv.org/abs/2304.11267
https://www.reddit.com/r/MachineLearning/comments/12zclus/d_google_researchers_achieve_performance/
本文來自微信公眾號：新智元 （ID：AI_era）
廣告聲明：文內(nèi)含有的對外跳轉(zhuǎn)鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時間，結(jié)果僅供參考，IT之家所有文章均包含本聲明。