Transformer 在大模型領域的地位可謂是難以撼動。不過，這個AI 大模型的主流架構在模型規模的擴展和需要處理的序列變長后，局限性也愈發凸顯了。Mamba的出現，正在強力改變著這一切。它優秀的性能立刻引爆了AI圈。

上周四， Vision Mamba（Vim）的提出已經展現了它成為視覺基礎模型的下一代骨干的巨大潛力。僅隔一天，中國科學院、華為、鵬城實驗室的研究人員提出了 VMamba：一種具有全局感受野、線性復雜度的視覺 Mamba 模型。這項工作標志著視覺 Mamba 模型 Swin 時刻的來臨。

• 論文標題：VMamba: Visual State Space Model
• 論文地址: https://arxiv.org/abs/2401.10166
• 代碼地址: https://github.com/MzeroMiko/VMamba

CNN 和視覺 Transformer（ViT）是當前最主流的兩類基礎視覺模型。盡管 CNN 具有線性復雜度，ViT 具有更為強大的數據擬合能力，然而代價是計算復雜較高。研究者認為 ViT 之所以擬合能力強，是因為其具有全局感受野和動態權重。受 Mamba 模型的啟發，研究者設計出一種在線性復雜度下同時具有這兩種優秀性質的模型，即 Visual State Space Model（VMamba）。大量的實驗證明，VMamba 在各種視覺任務中表現卓越。如下圖所示，VMamba-S 在 ImageNet-1K 上達到 83.5% 的正確率，比 Vim-S 高 3.2%，比 Swin-S 高 0.5%。

方法介紹

VMamba 成功的關鍵在于采用了 Selective Scan Space State Sequential Model（S6 模型）。該模型設計之初是用于解決自然語言處理（NLP）任務。與 ViT 中注意力機制不同，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 的整體框架與主流的視覺模型類似，其主要區別在于基本模塊（VSS block）中采用的算子不同。VSS block 采用了上述介紹的 2D-selective-scan 操作，即 SS2D。SS2D 保證了 VMamba 在線性復雜度的代價下實現全局感受野。

實驗結果

ImageNet 分類

通過對比實驗結果不難看出，在相似的參數量和 FLOPs 下:

• VMamba-T 取得了 82.2% 的性能，超過 RegNetY-4G 達 2.2%、DeiT-S 達 2.4%、Swin-T 達 0.9%。
• VMamba-S 取得了 83.5% 的性能，超過 RegNetY-8G 達 1.8%，Swin-S 達 0.5%。
• VMamba-B 取得了 83.2% 的性能（有 bug，正確結果將盡快在 Github 頁面更新），比 RegNetY 高 0.3%。

這些結果遠高于 Vision Mamba (Vim) 模型，充分驗證了 VMamba 的潛力。

COCO 目標檢測

在 COOCO 數據集上，VMamba 也保持卓越性能：在 fine-tune 12 epochs 的情況下，VMamba-T/S/B 分別達到 46.5%/48.2%/48.5% mAP，超過了 Swin-T/S/B 達 3.8%/3.6%/1.6% mAP，超過 ConvNeXt-T/S/B 達 2.3%/2.8%/1.5% mAP。這些結果驗證了 VMamba 在視覺下游實驗中完全 work，展示出了能平替主流基礎視覺模型的潛力。

ADE20K 語義分割

在 ADE20K 上，VMamba 也表現出卓越性能。VMamba-T 模型在 512 × 512 分辨率下實現 47.3% 的 mIoU，這個分數超越了所有競爭對手，包括 ResNet，DeiT，Swin 和 ConvNeXt。這種優勢在 VMamba-S/B 模型下依然能夠保持。

分析實驗

有效感受野

VMamba 具有全局的有效感受野，其他模型中只有 DeiT 具有這個特性。但是值得注意的是，DeiT 的代價是平方級的復雜度，而 VMamaba 是線性復雜度。

輸入尺度縮放

• 上圖（a）顯示，VMamba 在不同輸入圖像尺寸下展現出最穩定的性能（不微調）。有意思的是，隨著輸入尺寸從 224 × 224 增加到 384 × 384，只有 VMamba 表現出性能明顯上升的趨勢（VMamba-S 從 83.5% 上升到 84.0%），突顯了其對輸入圖像大小變化的穩健性。
• 上圖（b）顯示，VMamba 系列模型隨著輸入變大，復雜性呈線性增長，這與 CNN 模型是一致的。

最后，讓我們期待更多基于 Mamba 的視覺模型被提出，并列于 CNNs 和 ViTs，為基礎視覺模型提供第三種選擇。