ICLR2025|Spatial-Mamba：基於結構感知狀態融合的高效視覺狀態空間模型

2025-05-13 19:16 CVer

點選下方卡片，關注“CVer”公眾號

AI/CV重磅乾貨，第一時間送達
點選進入—>【頂會/頂刊】投稿交流群

新增微訊號：CVer2233，小助手會拉你進群！
掃描下方二維碼，加入CVer學術星球！可以獲得最新頂會/頂刊上的論文idea和CV從入門到精通資料，及最前沿應用！發論文/搞科研/漲薪，強烈推薦！

程式碼：github.com/EdwardChasel/Spatial-Mamba
論文：https://arxiv.org/abs/2410.15091

引言

狀態空間模型（State Space Models, SSMs）最初在控制理論、訊號處理等領域被廣泛使用，近年來，SSMs 被引入到深度學習中，特別是在自然語言處理（NLP）領域中表現出色。Mamba 模型透過選擇性機制和硬體感知最佳化，能夠高效地建模序列資料長距離依賴關係同時具有線性複雜度。
相比於一維序列資料，視覺資料本質上是二維空間結構的，直接將 SSMs 應用於影像面臨兩大難題：空間依賴性的丟失：現有方法（如圖 1 (a) – (c) ）透過不同掃描策略將影像展平為一維序列，難以保留影像空間結構資訊；計算效率低下：複雜的掃描路徑（如雙向、四向掃描）顯著增加了計算成本。
我們提出了一種新的視覺狀態空間模型 —— Spatial-Mamba，透過引入結構感知狀態融合（Structural-Aware State Fusion, SASF） 機制（如圖 1 (d) ），直接在狀態空間中建立領域連線，從而有效捕捉影像的空間依賴關係。我們深入分析了其與線性注意力和原始 Mamba 的關係，為理解自注意力機制在視覺表示學習中的作用提供了新的視角。Spatial-Mamba 不僅顯著提升了視覺上下文建模能力，還保持較高的計算效率。實驗結果表明，Spatial-Mamba 即使在單次掃描下，在影像分類、目標檢測和語義分割任務中均達到或超越了最先進的效能。

相關工作

狀態空間模型（SSMs）

Gu 等人首次將線性狀態空間層（LSSL）引入 HiPPO 框架，以高效處理長序列中的長程依賴關係。隨後，Gu 等人透過特殊矩陣引數化方式，顯著提高了 SSMs 的效率。S4 模型引發了一波結構化 SSMs 的研究浪潮。最近，Gu 和 Dao 提出了 Mamba 模型，將資料依賴的選擇性機制引入 S4 層，並以硬體友好的方式簡化了計算和架構，以線性複雜度實現了類似 Transformer 的建模能力。

視覺 SSMs

儘管傳統的 SSMs 在處理 NLP 序列資料和捕捉時間依賴關係方面表現優秀，但在處理視覺資料中固有的多維空間結構方面存在困難。最近的視覺 SSMs 傾向於設計多個掃描順序或模式以保持空間一致性，包括雙向、四向、連續、鋸齒形、基於視窗和基於拓撲的掃描。這些視覺 SSMs 已被用於多模態基礎模型、影像恢復、醫學影像分析等視覺任務，充分展示了 SSMs 在視覺資料理解中的潛力。

基礎概念

SSMs 通常用於分析序列資料並建模連續線性時不變（LTI）系統。輸入序列透過狀態變數轉換為輸出序列。整個動態系統可以透過線性狀態轉移方程和觀測方程來描述：
為了將連續時間 SSMs 有效地整合到深度學習框架中，必須對連續時間模型進行離散化，常用的技術是零階保持（ZOH）離散化，考慮離散時間步長，則引數，可以離散化為：
Mamba 指出真實現實場景往往不能被LTI系統準確地描述，而時變系統更多的關注相關資訊，能夠提供更真實的動態系統表示。Mamba 透過選擇性掃描機制自適應地感知上下文，狀態轉移引數依賴於輸入資料（，為簡單函式），從而嵌入上下文資訊。因此，離散後的時變狀態轉移和觀測方程如下：
圖 2 (a) 展示了上述過程的簡化示意圖。

Spatial-Mamba 方法

Spatial-Mamba 旨在捕捉潛在狀態空間中鄰域特徵的空間依賴關係。我們在原始 Mamba 狀態空間模型中引入了一種新的結構感知狀態融合（SASF）方程，Spatial-Mamba 的整個過程可以透過三個方程描述：狀態轉移、SASF 和觀測方程，其公式如下：
其中是原始狀態變數，是結構感知狀態變數，是鄰域集，是可學習權重，是位置t的第k個鄰域的索引。圖 2 (b) 展示了所提出的 Spatial-Mamba 中的 SSM 流程。與圖 2(a) 中的原始 Mamba 相比，原始狀態變數直接受其前一狀態的影響，而結構感知狀態變數透過融合機制合併了其他鄰近狀態變數的資訊。透過同時考慮全域性長距離和區域性空間資訊，融合狀態變數具備更豐富的上下文資訊，從而提高了對影像的適應性和理解能力。
具體而言，Spatial-Mamba 可以分三步實現。如圖 1(d) 所示，首先將輸入影像展平為 1D 序列，然後基於狀態轉換方程計算狀態變數，並將其轉換為 2D 格式。為了使每個狀態都能關注到 2D 空間中與之相鄰的狀態資訊，我們引入了 SASF 方程。對於狀態變數，我們使用權重 對鄰域中的相鄰狀態進行線性加權，從而有效地將區域性依賴關係整合到新狀態中。最後，透過觀測方程基於計算輸出。SASF 方法有助於模型在影像處理過程中融入區域性結構資訊，同時保留原始 Mamba 的優勢。
為了直觀地理解 SASF 機制，圖 3 可視化了原始狀態變數和結構感知狀態變數。可以看到，原始狀態變數（圖 3 (b) ）難以區分前景和背景。相反，經過 SASF 細化的結構感知狀態變數 （圖 3 (c)）有效地分離了這些區域。此外，圖 3 (d) 中的原始狀態變數僅顯示沿掃描方向的水平衰減（從左上角的最亮值逐漸變暗），而圖 3 (e) 中的融合狀態變數沿水平、垂直和對角線方向衰減。這種改進源於它能夠利用影像內的空間關係，從而實現更準確和更具上下文感知的特徵表示。

網路架構

Spatial-Mamba 的整體架構如圖 4 所示，由四個連續的模組組成，類似於 Swin-Transformer 的架構。

與原始Mamba和線性注意力的聯絡

我們深入分析了線性注意力、原始 Mamba 和 Spatial-Mamba 之間的相似點和差異點，從而更好地理解我們提出的方法。
  • 線性注意力：線性注意力是一種改進的自注意力（SA）機制，透過使用核函式將 SA 的計算複雜度降低為線性。對於輸入序列 ut，以及對應的查詢 qt、鍵 kt 和值 vt，省略歸一化的單頭線性注意力可以表示為：。令，則有 ，因此，線性注意力可以重寫為 。定義 和 ，則線性注意力可以表示成類似於 SSM 的形式：
  • Mamba：令初始狀態變數 為零並省略殘差項 ，定義，根據狀態轉移方程的遞推關係，可以將 Mamba 寫為下式：
  • Spatial-Mamba：類似地，Spatial-Mamba 可以表示為：
所有這三種正規化——線性注意力、Mamba 和 Spatial-Mamba —— 都可以表示為統一的矩陣乘法框架，即 。不同之處在於 的結構。對於線性注意力和 Mamba， 都為標準的下三角矩陣，而對於 Spatial-Mamba， 為鄰接矩陣。圖 5 提供了這些矩陣和對應啟用圖的視覺化結果。線上性注意力中，較亮值的位置在垂直方向上保持一致，表明 SA 機制偏向關注少量的影像標記。另一方面，Mamba 由於狀態轉移矩陣的影響，呈現出隨著時間推移的衰減趨勢，這種動態變換允許 Mamba 將其焦點轉移到先前的影像標記之間。與線性注意力和 Mamba 不同， Spatial-Mamba 考慮了更廣泛的空間鄰域內所有狀態，從而可以更全面地表示影像空間結構關係。圖 5 右側的啟用圖進一步表明，線性注意力關注的是有限區域，而 Mamba 則因其長距離上下文建模而捕獲了更廣闊的區域。Spatial-Mamba 不僅大大擴充套件了上下文建模的範圍，同時還實現了空間結構建模，即使彼此相距較遠，也能有效識別相關區域。

實驗結果

影像分類（ImageNet-1K）

我們在 ImageNet-1K 上評估了 Spatial-Mamba 的表徵學習能力。實驗結果表明，Spatial-Mamba-T 達到了83.5%的準確率，以更少引數量超越 CNN，Transformer 和 SSM 基線，其中較 VMamba-T 提升1%準確率。對於其他大小變體，Spatial-Mamba 也表現出顯著優勢。由於架構差異，Spatial-Mamba-T 比 VMamba-T 稍慢一些，但 Spatial-Mamba-S/B 的變體均比 VMamba 更快。此外，它們都比 CNN 和 Transformer 的方法快得多。

目標檢測與例項分割（COCO）

我們在 COCO 2017 資料集上評估了 Spatial-Mamba 在目標檢測和例項分割任務中的表現。實驗結果表明，Spatial-Mamba 以更少引數量和 FLOPs 達到更優的效能。

語義分割（ADE20K）

我們在 ADE20K 資料集上評估了 Spatial-Mamba 在語義分割任務中的表現。實驗結果表明，Spatial-Mamba 始終保持優異的效能。

消融實驗

我們在 ImageNet-1K 影像分類任務上對 Spatial-Mamba-T 的各個關鍵元件進行了消融實驗，包括鄰域集合、區域性增強、硬體最佳化和融合運算元等。實驗結果表明，Spatial-Mamba 的各個元件都對模型效能和效率有顯著影響。
更加詳細的實驗結果詳見我們的論文原文和 Github 主頁~

程式碼：github.com/EdwardChasel/Spatial-Mamba
論文：https://arxiv.org/abs/2410.15091

總結與展望

本文介紹了一種專為視覺任務設計的狀態空間模型 Spatial-Mamba。Spatial-Mamba 的關鍵在於提出的結構感知狀態融合 (SASF) 模組，該模組可有效捕獲影像空間依賴性，從而提高上下文建模能力。我們還深入分析了 Spatial-Mamba 與原始 Mamba 和線性注意力之間的關係，並將它們統一在同一個矩陣乘法框架下，為視覺表徵學習的自注意機制提供了更深入的理解。

參考文獻

[1] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In COLM, 2024.
[2] Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang, Qixiang Ye, and Yunfan Liu. Vmamba: Visual state space model. In NeurIPS, 2024.
[ *完整參考文獻列表詳見論文原文 ]

何愷明在MIT授課的課件PPT下載


在CVer公眾號後臺回覆：何愷明，即可下載本課程的所有566頁課件PPT！趕緊學起來！

ECCV 2024 論文和程式碼下載


在CVer公眾號後臺回覆：ECCV2024，即可下載ECCV 2024論文和程式碼開源的論文合集

CVPR 2024 論文和程式碼下載


在CVer公眾號後臺回覆：CVPR2024，即可下載CVPR 2024論文和程式碼開源的論文合集
CV垂直方向和論文投稿交流群成立
掃描下方二維碼，或者新增微訊號：CVer2233，即可新增CVer小助手微信，便可申請加入CVer-垂直方向和論文投稿微信交流群。另外其他垂直方向已涵蓋：目標檢測、影像分割、目標跟蹤、人臉檢測&識別、OCR、姿態估計、超解析度、SLAM、醫療影像、Re-ID、GAN、NAS、深度估計、自動駕駛、強化學習、車道線檢測、模型剪枝&壓縮、去噪、去霧、去雨、風格遷移、遙感影像、行為識別、影片理解、影像融合、影像檢索、論文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要備註：研究方向+地點+學校/公司+暱稱（如Mamba、多模態學習或者論文投稿+上海+上交+卡卡），根據格式備註，可更快被透過且邀請進群

▲掃碼或加微訊號: CVer2233，進交流群
CVer計算機視覺（知識星球）來了！想要了解最新最快最好的CV/DL/AI論文速遞、優質實戰專案、AI行業前沿、從入門到精通學習教程等資料，歡迎掃描下方二維碼，加入CVer計算機視覺（知識星球），已彙集上萬人！

▲掃碼加入星球學習
▲點選上方卡片，關注CVer公眾號
整理不易，請贊和在看


相關文章