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

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

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

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

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

論文概覽

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

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

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

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

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

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

DIFFUSSM 框架

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

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

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

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

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

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

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

實驗結(jié)果

類別條件圖像生成

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

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

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

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

無條件圖像生成

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

更多內(nèi)容請參考原論文。