來源：機器之心

替代注意力機制,，SSM 真的大有可為,？

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

眾所周知,，注意力機制是 Transformer 架構(gòu)的核心組件,，對于高質(zhì)量的文本、圖像生成都至關(guān)重要,。但它的缺陷也很明顯,，即計算復雜度會隨著序列長度的增加呈現(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 的未來充滿了信心。如今,，康奈爾大學和蘋果的這篇論文似乎又給 SSM 的應用前景增加了新的例證,。

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

論文概覽

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

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

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

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

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

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

DIFFUSSM 框架

作者的目標是設(shè)計一種擴散架構(gòu),，能夠在高分辨率下學習長程相互作用，而無需像分塊化那樣進行「長度縮減」,。與 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)在其他層組件（自適應層歸一化）中具有更多的參數(shù)。研究者在實驗中通過使用額外的 DIFFUSSM 層來匹配參數(shù),。

FLOPs：圖 3 比較了 DiT 和 DIFFUSSM 之間的 Gflops,。DIFFUSSM 一層的總 Flops 為

,，其中 α 代表 FFT 實現(xiàn)的常數(shù),。當 M = 2 且線性層主導計算時，這大約產(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é)果。

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

DIFFUSSM 在使用無分類器指導時的 FID 分數(shù)超越了所有模型，并在與 DiT 相比時保持了相當小的差距（0.01）,。需要注意的是,，在沒有應用無分類器指導的情況下，以減少 30% 的總 Gflops 訓練的 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 都有益處,。

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

無條件圖像生成

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

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