為了用更少的算力讓擴散模型生成高分辨率圖像，注意力機制可以不要，這是康奈爾大學和蘋果的一項最新研究所給出的結論。

眾所周知，注意力機制是 Transformer 架構的核心組件，對于高質量的文本、圖像生成都至關重要。但它的缺陷也很明顯，即計算復雜度會隨著序列長度的增加呈現二次方增長。這在長文本、高分辨率的圖像處理中都是一個令人頭疼的問題。

為了解決這個問題，這項新研究用一個可擴展性更強的狀態空間模型（SSM）主干替代了傳統架構中的注意力機制，開發出了一個名為 Diffusion State Space Model（DIFFUSSM）的新架構。這種新架構可以使用更少的算力，媲美或超越具有注意力模塊的現有擴散模型的圖像生成效果，出色地生成高分辨率圖像。

得益于上周「Mamba」的發布，狀態空間模型 SSM 正受到越來越多的關注。Mamba 的核心在于引入了一種新的架構 ——「選擇性狀態空間模型（ selective state space model）」，這使得 Mamba 在語言建模方面可以媲美甚至擊敗 Transformer。當時，論文作者 Albert Gu 表示，Mamba 的成功讓他對 SSM 的未來充滿了信心。如今，康奈爾大學和蘋果的這篇論文似乎又給 SSM 的應用前景增加了新的例證。

微軟首席研究工程師 Shital Shah 提醒說，注意力機制可能要從坐了很久的王座上被拉下來了。

論文概覽

圖像生成領域的迅速進展得益于去噪擴散概率模型（DDPMs）。這類模型將生成過程建模為迭代去噪潛變量，當執行足夠的去噪步驟時，它們能夠產生高保真度的樣本。DDPMs 捕捉復雜視覺分布的能力使其在推動高分辨率、照片級合成方面具有潛在的優勢。

在將 DDPMs 擴展到更高分辨率方面仍然存在重要的計算挑戰。主要瓶頸是在實現高保真生成時依賴自注意力。在 U-Nets 架構中，這個瓶頸來自將 ResNet 與注意力層相結合。DDPMs 超越了生成對抗網絡 (GANs)，但需要多頭注意力層。在 Transformer 架構中，注意力是中心組件，因此對于實現最新的圖像合成結果至關重要。在這兩種架構中，注意力的復雜性，與序列長度成二次方關系，所以當處理高分辨率圖像時將變得不可行。

計算成本促使以往的研究者們使用表示壓縮方法。高分辨率架構通常采用分塊化（patchifying）或多尺度分辨率。通過分塊化可以創建粗粒度表示，降低計算成本，但代價是犧牲關鍵的高頻空間信息和結構完整性。多尺度分辨率雖然可以減少注意層的計算，但也會通過降采樣減少空間細節并在應用上采樣時引入偽影。

擴散狀態空間模型（DIFFUSSM）是一種不使用注意力機制的擴散架構，它旨在解決在高分辨率圖像合成中應用注意力機制時出現的問題。DIFFUSSM 在擴散過程中采用了門控狀態空間模型 (SSM)。之前的研究表明，基于 SSM 的序列模型是一種有效而且高效的通用神經序列模型。通過使用這種架構，可以使 SSM 核心處理更細粒度的圖像表示，消除全局分塊化或多尺度層。為進一步提高效率，DIFFUSSM 在網絡的密集組件中采用沙漏 (hourglass) 架構。

作者在不同分辨率下驗證了 DIFFUSSM 的性能。在 ImageNet 上的實驗證明，在各種分辨率下，DIFFUSSM 在 FID、sFID 和 Inception Score 上都取得了一致的改進，并且總 Gflops 更少。

論文鏈接：https://arxiv.org/pdf/2311.18257.pdf

DIFFUSSM 框架

作者的目標是設計一種擴散架構，能夠在高分辨率下學習長程相互作用，而無需像分塊化那樣進行「長度縮減」。與 DiT 類似，該方法通過展平圖像并將其視為序列建模問題來實現。然而，與 Transformer 不同，這種方法在這個序列的長度上使用次二次（sub-quadratic）計算。

DIFFUSSM 的核心組件是優化處理長序列的門控雙向 SSM。為了提高效率，作者在 MLP 層中引入沙漏架構。這種設計在雙向 SSM 周圍交替擴展和收縮序列長度，同時在 MLP 中特定地減少序列長度。完整的模型架構如圖 2 所示。

具體來說，每個沙漏層接收經過縮短并展平的輸入序列 I ∈ R^(J×D)，其中 M = L/J 是縮小和放大的比例。同時，整個塊，包括雙向 SSM，在原始長度上進行計算，充分利用全局上下文。文中使用 σ 表示激活函數。對于 l ∈ {1 . . . L}，其中 j = ?l/M?，m = l mod M，D_m = 2D/M，計算方程如下所示：

作者在每個層中使用跳躍連接集成門控 SSM 塊。作者在每個位置集成了類標簽 y ∈ R^(L×1) 和時間步 t ∈ R^(L×1) 的組合，如圖 2 所示。

參數：DIFFUSSM 塊中參數的數量主要由線性變換 W 決定，其中包含 9D^2 + 2MD^2 個參數。當 M = 2 時，這產生了 13D^2 個參數。DiT 變換塊在其核心變換層中有 12D^2 個參數；然而，DiT 架構在其他層組件（自適應層歸一化）中具有更多的參數。研究者在實驗中通過使用額外的 DIFFUSSM 層來匹配參數。

FLOPs：圖 3 比較了 DiT 和 DIFFUSSM 之間的 Gflops。DIFFUSSM 一層的總 Flops 為，其中 α 代表 FFT 實現的常數。當 M = 2 且線性層主導計算時，這大約產生 7.5LD^2 Gflops。相比之下，如果在這個沙漏架構中使用全長的自注意力而不是 SSM，會有額外的 2DL^2 Flops。

考慮兩種實驗場景：1) D ≈ L = 1024，這將帶來額外的 2LD^2 Flops，2) 4D ≈ L = 4096，這將產生 8LD^2 Flops 并顯著增加成本。由于雙向 SSM 的核心成本相對于使用注意力的成本較小，因此使用沙漏架構對基于注意力的模型不起作用。正如前面討論的，DiT 通過使用分塊化來避免這些問題，以代價是壓縮表示。

實驗結果

類別條件圖像生成

表 1 是 DIFFUSSM 與目前所有的最先進的類別條件生成模型的比較結果。

當沒有使用無分類器指導時，DIFFUSSM 在 FID 和 sFID 兩方面均優于其他擴散模型，將之前非無分類器指導潛在擴散模型的最佳分數從 9.62 降至 9.07，同時使用的訓練步驟減少到原來的 1/3 左右。在訓練的總 Gflops 方面，未壓縮模型相較于 DiT 減少了 20% 的總 Gflops。當引入無分類器指導時，模型在所有基于 DDPM 的模型中獲得了最佳的 sFID 分數，超過了其他最先進的策略，表明 DIFFUSSM 生成的圖像對于空間失真更具魯棒性。

DIFFUSSM 在使用無分類器指導時的 FID 分數超越了所有模型，并在與 DiT 相比時保持了相當小的差距（0.01）。需要注意的是，在沒有應用無分類器指導的情況下，以減少 30% 的總 Gflops 訓練的 DIFFUSSM 已經超過了 DiT。U-ViT 是另一種基于 Transformer 的架構，但采用了基于 UNet 的架構，塊之間有長跳連接。U-ViT 在 256×256 分辨率下使用較少的 FLOPs，并在性能上表現更好，但在 512×512 數據集中情況并非如此。作者主要與 DiT 進行比較，為了公平，沒有采用這種長跳連接，作者認為采用 U-Vit 的思想可能對 DiT 和 DIFFUSSM 都有益處。

作者進一步在更高分辨率的基準上使用無分類器指導進行比較。DIFFUSSM 的結果相對強勁，并接近最先進的高分辨率模型，僅在 sFID 上不及 DiT，并獲得了可比較的 FID 分數。DIFFUSSM 在 302M 張圖像上進行了訓練，觀察了 40% 的圖像，使用的 Gflops 比 DiT 少了 25%。

無條件圖像生成

作者將模型的無條件圖像生成能力與現有基線進行比較。結果顯示在表 2 中。作者的研究發現，DIFFUSSM 在與 LDM 相當的訓練預算下取得了可比較的 FID 分數（差距為 - 0.08 和 0.07）。這個結果突顯了 DIFFUSSM 在不同基準和不同任務中的適用性。與 LDM 類似，由于只使用 ADM 總訓練預算的 25％，因此在 LSUN-Bedrooms 任務中，該方法并未超過 ADM。對于這個任務，最佳 GAN 模型在模型類別上勝過擴散模型。

更多內容請參考原論文。