Mamba帶火的SSM受到蘋果和康乃爾的關注：拋棄注意力分散模型

康乃爾大學與蘋果最新研究得出結論：為了用更少的算力產生高解析度影像，可以不使用注意力機制

眾所周知，注意力機制是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在網路的密集元件中採用了沙漏架構

#作者在不同解析度下驗證了 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 透過使用分塊化來避免這些問題，以代價是壓縮表示。

實驗結果

產生類別條件圖

下表是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在3.02億張圖像上進行了訓練，觀察了40%的圖像，使用的Gflops比DiT少了25%

無條件圖像生成

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

#請查閱原論文以取得更多詳細內容

#

以上是Mamba帶火的SSM受到蘋果和康乃爾的關注：拋棄注意力分散模型的詳細內容。更多資訊請關注PHP中文網其他相關文章！