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

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

• 論文標(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ù)擬合能力，然而代價(jià)是計(jì)算復(fù)雜較高。研究者認(rèn)為 ViT 之所以擬合能力強(qiáng)，是因?yàn)槠渚哂腥指惺芤昂蛣?dòng)態(tài)權(quán)重。受 Mamba 模型的啟發(fā)，研究者設(shè)計(jì)出一種在線性復(fù)雜度下同時(shí)具有這兩種優(yōu)秀性質(zhì)的模型，即 Visual State Space Model（VMamba）。大量的實(shí)驗(yàn)證明，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è)計(jì)之初是用于解決自然語言處理（NLP）任務(wù)。與 ViT 中注意力機(jī)制不同，S6 將 1D 向量中的每個(gè)元素（例如文本序列）與在此之前掃描過的信息進(jìn)行交互，從而有效地將二次復(fù)雜度降低到線性。

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

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

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

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

實(shí)驗(yàn)結(jié)果

ImageNet 分類

通過對(duì)比實(shí)驗(yàn)結(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) 模型，充分驗(yàn)證了 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é)果驗(yàn)證了 VMamba 在視覺下游實(shí)驗(yàn)中完全 work，展示出了能平替主流基礎(chǔ)視覺模型的潛力。

ADE20K 語義分割

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

分析實(shí)驗(yàn)

有效感受野

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

輸入尺度縮放

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

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