谷歌下場優化擴散模型，三星手機運行Stable Diffusion，12秒內出圖
Stable Diffusion 在圖像生成領域的知名度不亞于對話大模型中的 ChatGPT。其能夠在幾十秒內為任何給定的輸入文本創建逼真圖像。由于 Stable Diffusion 的參數量超過 10 億，并且由于設備上的計算和內存資源有限，因而這種模型主要運行在云端。
在沒有精心設計和實施的情況下，在設備上運行這些模型可能會導致延遲增加，這是由于迭代降噪過程和內存消耗過多造成的。
如何在設備端運行 Stable Diffusion 引起了大家的研究興趣，此前，有研究者開發了一個應用程序，該應用在 iPhone 14 Pro 上使用 Stable Diffusion 生成圖片僅需一分鐘，使用大約 2GiB 的應用內存。
此前蘋果也對此做了一些優化，他們在 iPhone、iPad、Mac 等設備上，半分鐘就能生成一張分辨率 512x512 的圖像。高通緊隨其后，在安卓手機端運行 Stable Diffusion v1.5 ，不到 15 秒生成分辨率 512x512 的圖像。
近日，谷歌發表的一篇論文中《 Speed Is All You Need: On-Device Acceleration of Large Diffusion Models via GPU-Aware Optimizations 》，他們實現了在 GPU 驅動的設備上運行 Stable Diffusion 1.4 ，達到 SOTA 推理延遲性能（在三星 S23 Ultra 上，通過 20 次迭代生成 512 × 512 的圖像僅需 11.5 秒）。此外，該研究不是只針對一種設備；相反，它是一種通用方法，適用于改進所有潛在擴散模型。
在沒有數據連接或云服務器的情況下，這項研究為在手機上本地運行生成 AI 開辟了許多可能性。Stable Diffusion 去年秋天才發布，今天已經可以塞進設備運行，可見這個領域發展速度有多快。
論文地址：https://arxiv.org/pdf/2304.11267.pdf
為了達到這一生成速度，谷歌提出了一些優化建議，下面我們看看谷歌是如何優化的。
方法介紹
該研究旨在提出優化方法來提高大型擴散模型文生圖的速度，其中針對 Stable Diffusion 提出一些優化建議，這些優化建議也適用于其他大型擴散模型。
首先來看一下 Stable Diffusion 的主要組成部分，包括：文本嵌入器（text embedder）、噪聲生成（noise generation）、去噪神經網絡（denoising neural network）和圖像解碼器（image decoder，如下圖 1 所示。
然后我們具體看一下該研究提出的三種優化方法。
專用內核：Group Norm 和 GELU
組歸一化（GN）方法的工作原理是將特征圖的通道（channel）劃分為更小的組，并獨立地對每個組進行歸一化，從而使 GN 對批大小的依賴性降低，更適合各種批大小和網絡架構。該研究沒有按順序執行 reshape、取均值、求方差、歸一化這些操作，而是設計了一個獨特的 GPU shader 形式的內核，它可以在一個 GPU 命令中執行所有這些操作，而無需任何中間張量（tensor）。
高斯誤差線性單元（GELU）作為常用的模型激活函數，包含大量數值計算，例如乘法、加法和高斯誤差函數。該研究用一個專用的 shader 來整合這些數值計算及其伴隨的 split 和乘法操作，使它們能夠在單個 AI 作畫調用中執行。
提高注意力模塊的效率
Stable Diffusion 中的文本到圖像 transformer 有助于對條件分布進行建模，這對于文本到圖像生成任務至關重要。然而，由于內存復雜性和時間復雜度，自 / 交叉注意力機制在處理長序列時遇到了困難。基于此，該研究提出兩種優化方法，以緩解計算瓶頸。
一方面，為了避免在大矩陣上執行整個 softmax 計算，該研究使用一個 GPU shader 來減少運算操作，大大減少了中間張量的內存占用和整體延遲，具體方法如下圖 2 所示。
另一方面，該研究采用 FlashAttention [7] 這種 IO 感知的精確注意力算法，使得高帶寬內存（HBM）的訪問次數少于標準注意力機制，提高了整體效率。
Winograd 卷積
Winograd 卷積將卷積運算轉換為一系列矩陣乘法。這種方法可以減少許多乘法運算，提高計算效率。但是，這樣一來也會增加內存消耗和數字錯誤，特別是在使用較大的 tile 時。
Stable Diffusion 的主干在很大程度上依賴于 3×3 卷積層，尤其是在圖像解碼器中，它們占了 90% 。該研究對這一現象進行了深入分析，以探索在 3 × 3 內核卷積上使用不同 tile 大小的 Winograd 的潛在好處。研究發現 4 × 4 的 tile 大小最佳，因為它在計算效率和內存利用率之間提供了最佳平衡。
該研究在各種設備上進行了基準測試：三星 S23 Ultra（Adreno 740）和 iPhone 14 Pro Max（A16）。基準測試結果如下表 1 所示：
很明顯，隨著每個優化被激活，延遲逐漸減少（可理解為生成圖像時間減少）。具體而言，與基線相比：在三星 S23 Ultra 延遲減少 52.2%；iPhone 14 Pro Max 延遲減少 32.9%。此外，該研究還對三星 S23 Ultra 端到端延遲進行評估，在 20 個去噪迭代 step 內，生成 512 × 512 像素圖像，不到 12 秒就達到 SOTA 結果。
小型設備可以運行自己的生成式人工智能模型，這對未來意味著什么？我們可以期待一波。
