MLNLP社群是國內外知名的機器學習與自然語言處理社群，受眾覆蓋國內外NLP碩博生、高校老師以及企業研究人員。

社群的願景是促進國內外自然語言處理，機器學習學術界、產業界和廣大愛好者之間的交流和進步，特別是初學者同學們的進步。

來源 | 極市平臺

作者丨科技猛獸

太長不看版

Mamba，一種使用類似 RNN 的狀態空間模型 (State Space Model, SSM) 作為 token mixer 的架構，能夠解決 Self-attention 的二次計算複雜度，並隨後引入視覺任務中。但是，與基於卷積或者 Attention 的模型相比，Mamba 在視覺方面的效能往往不足。

在本文中，作者深入研究了 Mamba 的本質，並得出結論：Mamba 非常適合具有長序列和自迴歸特徵的任務。視覺任務中的影像分類任務與上面的特徵都不匹配，因此作者假設 Mamba 在此任務中不是必需的；檢測和分割任務也不是自迴歸的，但它們遵循長序列特徵，因此作者認為探索 Mamba 在這些任務上的潛力仍然是值得的。

為了實證驗證這一假設，作者透過在移除核心 token mixer SSM 的同時堆疊 Mamba Block 來構建一系列名為 MambaOut 的模型。實驗結果有力地支援了我們的假設。具體來說，對於 ImageNet 影像分類，MambaOut 超過了所有視覺 Mamba 模型，這表明該任務確實沒有必要 Mamba。對於檢測和分割，MambaOut 無法與最先進的視覺 Mamba 模型的效能相匹配，展示了 Mamba 對長序列視覺任務的潛力。

圖1：左：Gated CNN 和 Mamba Block 的架構。Mamba 模組透過一個額外的狀態空間模型 (State Space Model, SSM) 來延展 Mamba。右：MambaOut 在 ImageNet 分類任務上超越了多個 Mamba 模型，包括 PlainMamba，Vision Mamba，VMamba

圖2：紀念科比布萊恩特—What can I say, Mamba out; 2016 年於 NBA farewell speech

本文做了哪些具體的工作

分析了 SSM 中類似 RNN 的機制，並在概念上得出結論，Mamba 適用於具有長序列和自迴歸特徵的任務。
檢查了視覺任務的特徵，認為 SSM 對於 ImageNet 影像分類任務是沒必要的，因為此任務不滿足自迴歸和長序列，但是探索 SSM 在檢測分割任務的潛力仍然很有價值，因為這些任務儘管不是自迴歸的，但是符合長序列的特徵。
開發了一系列名為 MambaOut 的模型，該模型基於 Gated CNN Block，但沒有 SSM。實驗表明，MambaOut 在 ImageNet 影像分類中有效地超過了視覺 Mamba 模型。

MambaOut：在視覺任務中，我們真的需要 Mamba 嗎？

論文名稱：MambaOut: Do We Really Need Mamba for Vision? (Arxiv 2024.05)

論文地址：

http://arxiv.org/pdf/2405.07992

程式碼連結：

http://github.com/yuweihao/MambaOut

1.1 在視覺任務中，我們真的需要 Mamba 嗎？

近年來，Transformer[1]已經成為各種任務的主流骨幹，支援 BERT、GPT 系列和 ViT 等突出模型。然而，Transformer 的 token mixer Self-attention 有關於序列長度的二次計算複雜度，對長序列任務提出了重大的挑戰。為了解決這個問題，前人的工作引入了各種各樣的具有線性複雜度的 token mixer，比如動態卷積[2]、Linformer[3]、Longformer[4]、Big Bird[5]和 Performer[6]等。最近，出現了一種新的類似 RNN 的模型浪潮[7][8][9]，因其可並行化訓練以及對長序列執行有效推理的能力，引起了社群的極大興趣。值得注意的是，RWKV[8]和 Mamba[9]等模型被證明對大語言模型 (LLM) 的主幹[10]有效。

受這些類似 RNN 的模型展示出的 promising 的能力的啟發，各種研究工作試圖將 Mamba[9]引入到視覺工作中，例如 Vision Mamba[11]，VMamba[12]，LocalMamba[13]和 PlainMamba[14]等工作。本著 RNN 的精神，Mamba 的令牌混合器是結構化狀態空間模型 (State Space Model, SSM)。然而，他們的實驗表明，與最先進的基於卷積[15][16]和基於 Attention[17][18]的模型相比，基於 SSM 的視覺模型的效能不足。這就產生了一個研究問題：在視覺任務中，我們真的是否需要 Mamba？

1.2 本文有哪些新的發現？

在本文中，作者研究了 Mamba 的性質，並在概念上總結了 Mamba 非常適合具有兩個關鍵特徵的任務：長序列 (long-sequence) 和**自迴歸 (autoregressive)**，原因是 SSM 固有的 RNN 機制[19][20]。但不是所有的視覺任務都具有這兩個特徵。例如，ImageNet 影像分類兩者都不符合，COCO 目標檢測和例項分割，ADE20K 語義分割任務僅符合長序列這個特徵。另一方面，自迴歸特性要求每個 token 僅從前一個 token 和當前 token 中聚合資訊，也就是一種 causal mode 的 token mixing 策略。事實上，所有的視覺識別任務都屬於是理解領域而非生成任務，意味著模型一次可以看到整個影像。因此，在視覺識別模型中對 token mixing 施加額外的 causal 約束可能會導致效能下降。雖然這個問題可以透過新增 Bidirectional 分支[21]來緩解，但這個問題不可避免地存在於每個分支中。

基於上述概念討論，作者提出以下兩個假設：

假設1：影像分類任務沒必要使用 SSM。因為這個任務既不屬於長序列任務，也不屬於自迴歸任務。
假設2：SSM 可能對目標檢測，例項分割和語義分割任務有益，因為它們儘管不是自迴歸的，但遵循長序列特徵。

為了驗證這兩個假設，作者開發了一系列模型，稱為 MambaOut，透過堆疊 Gated CNN[22]Block。Gated CNN 和 Mamba Block 之間的主要區別在於 SSM 的存在，如圖1所示。實驗結果表明，更簡單的 MambaOut 模型已經超過了視覺 Mamba 模型的效能，這反過來又驗證了本文假設1。作者還透過實驗證明，MambaOut 在檢測和分割任務中沒有達到最先進的視覺 Mamba 模型的效能，表明 SSM 在這些任務上的潛力，並有效地驗證了本文假設2。

1.3 Mamba 適合什麼任務？

在本節中，作者首先討論 Mamba 模型適合的任務的特徵。接下來，檢查視覺識別任務是否符合這些特徵。基於檢查結果，作者提出了關於 Mamba 在視覺的必要性的假設。

Mamba 模型的 token mixer 是 selective SSM[9][20]，其包含4個 input-dependent 的引數, 並將它們透過下式變為 :

然後，SSM 的序列到序列變換可以表示為：

其中表示時間步長, 表示輸入, 表示隱藏狀態, 表示輸出, 上式2的迴圈屬性將 SSM 與 Causal Attention 區分開來。隱藏狀態可以看作是儲存所有歷史資訊的固定大小的 memory。透過式2, 更新此記憶體, 同時其大小保持不變。固定大小意味著 memory 不可避免地是有損的, 但它確保將 memory 與當前輸入整合的計算複雜度保持不變。相反, Causal Attention 將來自先前 token 的所有鍵和值儲存為其 memory, 它透過為每個新輸入添加當前 token 的鍵和值來擴充套件。這個 memory 在理論上是無損的。然而, 隨著輸入更多的 tokens, memory 的大小也會增加, 從而增加了將 memory 與當前輸入整合的複雜性。圖2進一步說明了類似 RNN 的模型和 Causal Attention 之間的記憶機制的差異。

圖3：從 memory 的角度說明 Causal attention 和 RNN-like 模型的區別，x_i 代表第 i 步的輸入 tokens。(a): Causal attention 將之前的所有 key 和 value 存入 memory。memory 是無損的；(b)：RNN-like 模型將之前的 tokens 壓縮為固定大小的狀態，存入 memory。memory 是有損的

SSM 的 memory 在本質上是有損的，它在邏輯上落後於 Attention 的無損 memory。因此，Mamba 無法展示處理短序列方面的優勢。而 Attention 在短序列處理上很容易處理好。然而，在涉及長序列的場景中，Attention 會因其二次複雜度而發生變化。在這種情況下，Mamba 可以突顯其在將記憶體與當前輸入合併時的效率，從而平滑地處理長序列。因此，Mamba 特別適合處理長序列。

雖然 SSM 的迴圈性質 (式2) 允許 Mamba 有效地處理長序列，但它引入了一個顯著的侷限性：只能訪問前一個時間步和當前時間步的資訊。如圖3所示，這種型別的 token mixing 稱為 Causal Mode，可以表示為：

其中和分別表示第個 token 的輸入和輸出。由於其因果性質，這種模式非常適合自迴歸生成任務。

另一種模式稱為完全可見模式，其中每個 token 都可以聚合來自所有先前和後續 tokens 的資訊。這意味著每個 token 的輸出取決於所有 tokens 的輸入：

式中，代表總的 token 數。完全可見模式適用於理解任務，其中模型可以同時訪問所有輸入。

預設情況下，Attention 處於完全可見的模式，但它可以透過將 Causal Mask 應用於注意力圖輕鬆轉變為 Causal Mode。RNN-like 模型由於其迴圈特性固有地以因果模式執行。由於這種固有特性，RNN-like 模型無法轉換為完全可見的模式。儘管 RNN 可以使用雙向分支近似完全可見的模式，但每個分支仍然是 Causal Mode。因此，由於 Mamba 的迴圈特性的固有限制，導致它非常適合需要 Causal token mixing 的任務。

總之，Mamba 非常適合顯示以下特徵的任務：

特徵1：該任務涉及處理長序列。
特徵2：該任務需要 Causal 的 token mixing 模式。

接下來，作者將討論視覺識別任務是否表現出這兩個特徵。

1.4 視覺任務有長序列嗎？

作者探討了視覺識別任務是否需要長序列建模。作者使用 Transformer 作為案例研究。考慮一個常見的 MLP 比為 4 的 Transformer Block; 假設其輸入的長度為 , Embedding dimension 的維度為 , Block 的 FLOPs 可以計算為:

由此，推匯出二次項與線性項的比率為：

如果，則二次項的計算負載超過了線性項的計算量。這提供了一個簡單的指標來確定任務是否涉及長序列。例如, 在 ViT-S 中有384個 channel, 閾值 , ViT-B 中有 768 個 channel, 閾值。

對於 ImageNet 影像分類任務, 典型的輸入影像大小為 , 就會有個 token, Patch Size 為。顯然, 196 遠小於和 , 這表明 ImageNet 影像分類任務不屬於長序列任務。

對於目標檢測和例項分割任務, 推理影像大小為 , 在 ADE20K 語義分割任務, 推理影像大小為 , 給定 Patch Size 為 , token 的數量約為。由於以及 , 因此檢測和 ADE20K 語義分割都可以被認為是長序列任務。

1.5 視覺任務需要 Causal 的 token mixing 模式嗎？

圖4：(a) 兩種模式的 token mixing 策略。共計 T 個 tokens，全部可見模式允許 token t 根據所有的 tokens 來計算輸出，而 Causal 模式則限制 token t 根據當前以及之前的所有的 tokens 來計算輸出。(b) 作者將完全可見模式修改為 Causal 模式，並觀察到 ImageNet 的效能下降

如上圖3所示，完全可見模式進行 token mixing 時候不限制範圍，而 Causal 模式限制僅僅訪問之前 token 的資訊。視覺識別被歸類為理解任務，其中模型一次可以看到整個影像，消除了對 token mixing 的限制。對 token mixing 施加額外的約束可能會降低模型效能。如圖 3(b) 所示，當 Causal 的限制應用於 Vision Transformers (ViT) 時，可以觀察到效能的顯著下降。一般來說，完全可見的模式適用於理解任務，而因果模式更適合自迴歸任務。這也是為什麼 BERT 和 ViT 比 GPT-1/2 和 Image GPT 更多用於理解任務。

1.6 關於 Mamba 在視覺的必要性的假設

基於之前的討論，作者總結了假設，即為視覺識別任務引入 Mamba 的必要性：

假設 1：沒有必要在 ImageNet 上引入 SSM 進行影像分類，因為此任務不滿足特徵1或特徵2。
假設 2：仍然值得進一步探索 SSM 在檢測和分割方面的潛力，因為這些任務與特徵1一致，儘管不滿足特徵2。

1.7 Gated CNN 和 MambaOut

接下來, 作者透過實驗來驗證之前的假設。如圖 1(a) 所示, Mamba Block 基於 Gated CNN Block 。Gated CNN 和 Mamba 的元架構都可以被認為是 MetaFormer 中的 token mixer 和 MLP 的簡化整合, 類似於 MetaNeXt 。形式上, 給定輸入 , 元架構表示為:

式中， TokenMixer()代表 token mixer，其在 Gated CNN 和 Mamba 中分別為：

可以發現，Gated CNN 和 Mamba Block 之間的主要區別在於 SSM 的存在。作者開發了一系列模型，稱為 MambaOut，它基於沒有 SSM 的 Gated CNN Block。透過 MambaOut 來評估 Mamba 對視覺識別任務的必要性。

具體而言，作者按照 ConvNeXt 的做法將 Gated CNN 的 token mixer 指定為 7×7 的 Depth-Wise Convolution。此外，為了提高實際速度，我們僅對部分 channels 進行 Depth-Wise Convolution。如下圖5所示，Gated CNN 的實現簡單而優雅。與 ResNet 類似，我們採用 4 階段框架，透過在每個階段堆疊 Gated CNN 的來構建 MambaOut，如圖6所示。

圖6：(a) MambaOut 的總體架構；(b) Gated CNN Block 的設計

1.8 ImageNet 實驗結果

ImageNet 實驗結果如下圖7所示。本文不包含 SSM 的 MambaOut 模型始終優於包含所有模型大小的 SSM 的視覺 Mamba 模型。比如，MambaOut-Small 模型實現了 84.1% 的 top-1 精度，比 LocalVMamba-S 高 0.4%，同時只需要 79% 的 MAC。這些結果有力地支援了本文假設 1，即沒有必要在 ImageNet 上引入 SSM 進行影像分類，這與 Occam 的 razor 的原理對齊。

此外，與最先進的卷積和注意力模型相比，視覺 Mamba 模型目前表現出顯著的效能差距。例如，CAFormer-M36[18]採用了簡單的可分離卷積和標準注意機制等傳統 token mixer，在準確率上優於所有大小相當的視覺 Mamba 模型。

1.9 COCO 目標檢測和例項分割實驗結果

作者使用 Mask R-CNN[27]作為目標檢測和例項分割框架。雖然 MambaOut 在 COCO 目標檢測和例項分割任務中可以超越一些視覺 Mamba 模型，但它仍然落後於最先進的視覺 Mamba，例如 VMamba 和 LocalVMamba。這種效能差異強調了在長序列視覺任務中整合 Mamba 的好處，驗證了本文假設2。然而，與最先進的卷積注意混合模型 TransNeXt[28]相比，視覺 Mamba 仍然表現出顯著的效能差距。視覺 Mamba 需要透過在視覺檢測任務中優於其他最先進的模型來進一步驗證其有效性。

1.10 ADE20K 語義分割實驗結果

作者使用 UperNet[29]作為語義分割框架。ADE20K 語義分割的效能趨勢類似於 COCO 目標檢測的結果。MambaOut 可以勝過一些視覺 Mamba 模型，但不能匹配最先進的 Mamba 模型的結果。例如，LocalVMamba-T 在單尺度 (SS) 和多尺度 (MS) 評估中都超過了 MambaOut-Tiny 0.5 mIoU，進一步驗證實了本文假設2。此外，與包含了卷積和注意力機制的更先進的混合模型 (如 SG-Former[30]和 TransNeXt[28]) 相比，視覺 Mamba 表現出了顯著的效能缺陷。因此，還應該透過在分割任務中提供更強的效能來進一步展示其長序列建模的優勢。

參考

^Attention is All you Need
^Lite Transformer with Long-Short Range Attention
^Linformer: Self-Attention with Linear Complexity
^Longformer: The Long-Document Transformer
^Big Bird: Transformers for Longer Sequences
^Masked Language Modeling for Proteins via Linearly Scalable Long-Context Transformers
^Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
^abRWKV: Reinventing RNNs for the Transformer Era
^abcdMamba: Linear-Time Sequence Modeling with Selective State Spaces
^Jamba: A hybrid transformer-mamba language model
^Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model
^VMamba: Visual State Space Model
^LocalMamba: Visual State Space Model with Windowed Selective Scan
^Plainmamba: Improving non-hierarchical mamba in visual recognition
^A ConvNet for the 2020s
^Scaling Up Your Kernels to 31×31: Revisiting Large Kernel Design in CNNs
^Going deeper with Image Transformers
^abMetaFormer Baselines for Vision
^Combining Recurrent, Convolutional, and Continuous-time Models with Linear State-Space Layers
^abEfficiently Modeling Long Sequences with Structured State Spaces
^Bidirectional recurrent neural networks
^Language Modeling with Gated Convolutional Networks
^Long Short-Term Memory
^Language Modeling with Gated Convolutional Networks
^MetaFormer Is Actually What You Need for Vision
^InceptionNeXt: When Inception Meets ConvNeXt
^Mask R-CNN
^abTransNeXt: Robust Foveal Visual Perception for Vision Transformers
^Unified perceptual parsing for scene understanding
^SG-Former: Self-guided Transformer with Evolving Token Reallocation