Transformer 在大模型領(lǐng)域的地位可謂是難以撼動。不過，這個AI 大模型的主流架構(gòu)在模型規(guī)模的擴(kuò)展和需要處理的序列變長后，局限性也愈發(fā)凸顯了。Mamba的出現(xiàn)，正在強(qiáng)力改變著這一切。它優(yōu)秀的性能立刻引爆了AI圈。

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

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

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

方法介紹

VMamba 成功的關(guān)鍵在于采用了 Selective Scan Space State Sequential Model（S6 模型）。該模型設(shè)計之初是用于解決自然語言處理（NLP）任務(wù)。與 ViT 中注意力機(jī)制不同，S6 將 1D 向量中的每個元素（例如文本序列）與在此之前掃描過的信息進(jìn)行交互，從而有效地將二次復(fù)雜度降低到線性。

然而，由于視覺信號（如圖像）不像文本序列那樣具有天然的有序性，因此無法在視覺信號上簡單地對 S6 中的數(shù)據(jù)掃描方法進(jìn)行直接應(yīng)用。為此研究者設(shè)計了 Cross-Scan 掃描機(jī)制。Cross-Scan 模塊（CSM）采用四向掃描策略，即從特征圖的四個角同時掃描（見上圖）。該策略確保特征中的每個元素都以不同方向從所有其他位置整合信息，從而形成全局感受野，又不增加線性計算復(fù)雜度。

在 CSM 的基礎(chǔ)上，作者設(shè)計了 2D-selective-scan（SS2D）模塊。如上圖所示，SS2D 包含了三個步驟：

• scan expand 將一個 2D 特征沿 4 個不同方向（左上、右下、左下、右上）展平為 1D 向量。
• S6 block 獨立地將上步得到的 4 個 1D 向量送入 S6 操作。
• scan merge 將得到的 4 個 1D 向量融合為一個 2D 特征輸出。

上圖為本文提出的 VMamba 結(jié)構(gòu)圖。VMamba 的整體框架與主流的視覺模型類似，其主要區(qū)別在于基本模塊（VSS block）中采用的算子不同。VSS block 采用了上述介紹的 2D-selective-scan 操作，即 SS2D。SS2D 保證了 VMamba 在線性復(fù)雜度的代價下實現(xiàn)全局感受野。

實驗結(jié)果

ImageNet 分類

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

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

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

COCO 目標(biāo)檢測

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

ADE20K 語義分割

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

分析實驗

有效感受野

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

輸入尺度縮放

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

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