VMamba: Visual State Space Model

論文地址: https://arxiv.org/abs/2401.10166
代碼地址: https://github.com/MzeroMiko/VMamba

摘要

• 卷積神經網絡(CNN)和視覺Transformer(ViT)是視覺表示學習的兩種最流行的基礎模型
• CNN表現出卓越的可擴展性和線性復雜度。
• ViT在圖像分辨率方面超過了CNN, 但是復雜性確實二次方。Vit通過結合全局感受野和動態權重實現了卓越的視覺建模性能。
• 作者在繼承上述組件的情況下引入了狀態空間模型，提出了視覺狀態空間模型(VMamba)。該模型在不犧牲全局感受野和動態權重的情況下實現了線性復雜度。
• 通過引入交叉掃描模塊(CSM)解決方向敏感問題。
• 實驗證明，VMamba在各種視覺任務表現良好，并且隨著圖像分辨率的提高而表現出更明顯的優勢。

引言

• 視覺表示學習是計算機視覺領域最基本的研究課題之一。深度基礎模型主要分為卷積神經網絡(CNN)和Vision Transformer(ViTs)兩個主要類別。

• 然而，注意力機制在圖像大小方面需要二次方的 復雜度，這在處理下游密集預測任務時計算開銷昂貴。為了解決這個問題，人們投入大量的經歷通過限制計算窗口的大小或步幅來提高注意力的效率。但是這種是以限制感受野的規模為代價的。因此，作者設計了一種具有線性復雜度的心得視覺基礎模型，同時仍然保留與全局感受野和動態權重相關的優勢。

• 引入了視覺狀態空間模型(VMamba), 用于高效的視覺表示學習。它通過繼承自選擇性掃描性狀態序列模型(S6)。但是因為無法估計和未知掃描補丁的關系，直接使用這種策略會導致接受域受限。作者稱這種問題為"方向敏感"問題，并引入交叉掃描模塊(CSM)來解決。CSM采用四向掃描策略，即從整個特征的四個角開始掃描映射到相反的為止。如下圖所示。該策略保持特征圖中的每個元素集成自不同方向的所有其他位置的信息，從而在不增加線性計算復雜度的情況下呈現全局感受野。
  

• VMamba的Tiny、Small、Base分別包含22M、44M、75M參數量。能夠在FLOPs穩步增加的情況下獲得與ViT相當的性能。

• 貢獻:

  • 提出了VMamba, 一種具有全局感受野和動態權重的視覺狀態空間模型，用于視覺表示學習。
  • 引入交叉掃描模塊(CMS)是為了彌補一維陣列掃描和二維平面掃描之間的差距，促進S6在不影響接受范圍的情況下擴展到視覺數據
  • 證明VMamba在圖像分類、目標檢測和語義分割等各種視覺任務具有強大的潛力。

相關工作

具體內容略，主要為卷積神經網絡(CNN)、Vision Transformer(ViTs)和狀態空間模型(SSMs), 狀態空間模型詳解可參考Mamba

Preliminaries

狀態空間模型與相關工作一樣可參考Mamba
連續時間的SSM可以表示為線性常微分方程(ODEs, linear ordinary differential equations):
$$\begin{array}{rl}{h}^{\prime }(t)& =\mathbf{A}h(t)+\mathbf{B}\mu (t),\\ y(t)& =\mathbf{C}h(t)+D\mu (t),\end{array}$$
其中, $\mathbf{A}\in {\mathbb{R}}^{N\times N},\mathbf{B}\in {\mathbb{R}}^{N\times 1},\mathbf{C}\in {\mathbb{R}}^1$是權重參數

SSM的離散化: 狀態空間模型(SSMs)作為連續時間模型，需要進行離散化才能更好的集成到深度學習算法中。可以參照Mamba中S4的部分。

選擇掃描機制: 為了解決LTI SSMs(SSMs原始公式)獲取上下文信息的局限性，Gu等人提出了一種新的SSMs參數化方法, 該方法集成了一個輸入依賴的選擇機制(S6)。然而，在選擇性SSMs的情況下，卷積不適應動態權重導致時變加權參數對隱藏狀態難以進行有效計算。而離散化可以使用線性復雜度的關聯掃描算法有效地計算出$y_b$。

方法

網絡結構

作者在三個尺度上開發了VMamba: VMamba-Tiny(VMamba-T)、VMamba-Small(VMamba-S)和VMamba-Base(VMamba-B)。VMamba-T架構的概述如下圖中的a所示。首先將輸入圖像$I\in {\mathbb{R}}^{H\times W\times 3}$分割成多個patch, 得到一個維度為$H/4\times W/4$的2d的特征圖空間。隨后，使用多個網絡階段創建分辨率為$H/8\times W/8,H/16\times W/16,H/32\times W/32$的分層表示。每個階段都包括一個下采樣層(第一階段除外)，然后是一個堆疊的視覺狀態空間(VSS)塊。

VSS塊是Mamba塊(上圖的b)的視覺對應部分。新提出的2d選擇掃描模塊(SS2D)替換VSS塊的初始架構(上圖的c)作為Mamba的核心的同時實現全局接受域，動態權重(即選擇性)和線性復雜度。

為了進一步提高計算效率，消除了整個乘法分支(上圖c中的紅框部分)，因為門控機制的影響是通過SS2D的選擇性來說實現的。因此，生成的VSS塊(上圖的d)由一個包含兩個殘差模塊的單一網絡分支組成，模仿了一個普通transformer的架構。

2D-Selective-Scan for Vision Data(SS2D)

S6中的掃描操作的順序性質難以用于視覺數據。因為視覺數據本質上是非順序的，包含空間信息(如局部紋理和全局結構)。為了解決該問題, S4ND用卷積運算重新指定了SSM, 通過外積將內核從一維擴展到二維，然而這種修改使得權重無法獨立于輸入，從而限制了捕獲上下文信息的能力。因此，作者基于選擇性掃描方法提出了了2D選擇性掃描模塊(SS2D)，在不影響其優勢的情況下使S6適應視覺數據。

如上圖所示，SS2D中的數據前向包含三個步驟: 交叉掃描、S6的選擇性掃描和交叉合并。給定輸入數據，SS2D首先沿著輸入patch沿著四個不同的遍歷路徑展開成序列(即交叉掃描)，使用單獨的S6塊并行的處理每個patch序列，然后對合成序列進行重構和合并, 以形成輸出映射(即交叉合并)。通過采用互補的以為遍歷路徑，SS2D使圖像中的每個像素能夠有效地整合其他所有像素的信息，從而促進在二維空間中建立全局接受域的過程。