Mamba 成功的關鍵在于采用了 Selective Scan Space State Sequential Model(S6 模型)。是用于解決自然語言處理(NLP)任務。與 transformer中注意力機制不同,Mamba的S6 將 1D 向量中的每個元素(例如文本序列)與在此之前掃描過的信息進行交互,從而有效地將二次復雜度降低到線性。
然而,由于視覺信號(如圖像)不像文本序列那樣具有天然的有序性,因此無法在視覺信號上簡單地對 S6 中的數據掃描方法進行直接應用。為此研究者設計了 Cross-Scan 掃描機制。Cross-Scan 模塊(CSM)采用四向掃描策略,即從特征圖的四個角同時掃描(見上圖)。該策略確保特征中的每個元素都以不同方向從所有其他位置整合信息,從而形成全局感受野,又不增加線性計算復雜度。
在 CSM 的基礎上,作者設計了 2D-selective-scan(SS2D)模塊。如上圖所示,SS2D 包含了三個步驟:
-
scan expand 將一個 2D 特征沿 4 個不同方向(左上、右下、左下、右上)展平為 1D 向量。
-
S6 block 獨立地將上步得到的 4 個 1D 向量送入 S6 操作。
-
scan merge 將得到的 4 個 1D 向量融合為一個 2D 特征輸出。
上圖為本文提出的 VMamba 結構圖。VMamba 的整體框架與主流的視覺模型類似,如上圖?(b)所示。經過Layer Normalization后,輸入被分成兩個分支。在第一個分支中,輸入經過一個線性層,然后是一個激活函數。在第二個分支中,輸入通過線性層、深度可分離卷積和激活函數進行處理,然后輸入到2D選擇性掃描(SS2D)模塊中進行進一步的特征提取。隨后,使用Layer Normalization對特征進行歸一化,然后使用第一個分支的輸出執行逐元素的生成,以合并兩條路徑。最后,使用線性層混合特征,并將此結果與殘差連接相結合,形成VSS塊的輸出。本文默認采用SiLU作為激活函數。
????????其主要區別在于基本模塊(VSS block)中采用的算子不同。VSS block 采用了上述介紹的 2D-selective-scan 操作,即 SS2D。SS2D 保證了 VMamba 在線性復雜度的代價下實現全局感受野。
SS2D由三個部分組成:掃描expanding操作、S6塊操作和掃描merging操作。如圖2(a)所示,掃描expanding操作沿著四個不同的方向(左上到右下、左下到右上、右下到左上、右上到左下)將輸入圖像展開成序列。然后通過S6塊對這些序列進行特征提取,確保各個方向的信息被徹底掃描,從而捕獲不同的特征。隨后,如圖2(b)所示,掃描merging操作將來自四個方向的序列相加并合并,將輸出圖像恢復為與輸入相同的大小。源自Mamba[16]的S6塊在S4[17]之上引入了一種選擇機制,通過根據輸入調整SSM的參數。這使模型能夠區分并保留相關信息,同時過濾掉不相關的信息。