來源：機(jī)器之心

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

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

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

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

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

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

論文概覽

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

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

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

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

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

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

DIFFUSSM 框架

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

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

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

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

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

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

,，其中 α 代表 FFT 實(shí)現(xiàn)的常數(shù),。當(dāng) M = 2 且線性層主導(dǎo)計(jì)算時(shí)，這大約產(chǎn)生 7.5LD^2 Gflops,。相比之下,，如果在這個(gè)沙漏架構(gòu)中使用全長的自注意力而不是 SSM，會(huì)有額外的 2DL^2 Flops,。

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

實(shí)驗(yàn)結(jié)果

類別條件圖像生成

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

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

DIFFUSSM 在使用無分類器指導(dǎo)時(shí)的 FID 分?jǐn)?shù)超越了所有模型,，并在與 DiT 相比時(shí)保持了相當(dāng)小的差距（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 進(jìn)行比較,，為了公平，沒有采用這種長跳連接,，作者認(rèn)為采用 U-Vit 的思想可能對(duì) DiT 和 DIFFUSSM 都有益處,。

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

無條件圖像生成

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

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