正如標題所言「Attention is all you need」，Transformer已成為不同領域的「霸主」，包括計算機視覺、自然語言處理、語音處理和機器人技術。

第一個挑戰Transformer架構的是Mamba，一種新的狀態空間模型（SSM），它具有線性時間復雜度，并在多個語言建模任務中超越或與Transformer媲美。

但在不同的視覺任務上，Vision Transformer (ViT) 和卷積神經網絡 (CNN) 架構的骨干網絡，仍然優于基于Mamba的視覺模型。

而這一次，英偉達高級工程師Ali Hatamizade，宣布被頂會CVPR-2025接受的MambaVision，在視覺任務上超越以往的模型，而設計的關鍵在于將Mamba和Transformer混合。

正如圖1所示，在ImageNet-1K基準上，MambaVision的Top-1準確率和圖像處理能力達到了新的Pareto最優點，超越了Mamba、CNN和ViT基于的模型，有時差距非常顯著。

在下游任務如目標檢測、實例分割以及語義分割中，采用MambaVision作為骨干網絡的模型在MS COCO和ADE20數據集上分別超越了同等規模的對比模型。

MambaVision是首個針對計算機視覺應用，結合Mamba和Transformer的混合架構的嘗試。主要貢獻總結如下：

1 引入了重新設計的適用于視覺任務的Mamba模塊，提升了相較于原始Mamba架構的準確性和圖像處理能力。

2 系統性地研究了Mamba和Transformer模塊的融合模式，并展示了在最終階段加入自注意力模塊，顯著提高了模型捕捉全局上下文和長距離空間依賴的能力。

論文鏈接:https://arxiv.org/abs/2407.08083

在這項工作中，作者系統地重新設計了Mamba模塊，使其更加適合視覺任務。

新方法是一種混合架構，結合了新提出的公式（即MambaVision Mixer和MLP）以及Transformer模塊。

具體來說，研究了不同的集成模式，比如以等參數方式將Transformer模塊添加到早期、中間和最終層，或者每隔l層添加一次。

分析表明，在最終階段利用多個自注意力模塊，可以顯著增強捕捉全局上下文和長程空間依賴的能力。

使用混合架構相較于純Mamba或ViT模型，圖像處理能力也得到了顯著提升。

網絡架構

宏觀架構

如圖2所示，MambaVision 采用了分層架構，由4個不同的階段組成。

前兩個階段使用基于CNN的層，負責在較高輸入分辨率下進行快速特征提取，而第3和第4階段則包括了新提出的 MambaVision和Transformer模塊。

具體來說，給定一個大小為H×W×3的圖像，輸入首先被轉換為大小為H/4×W/4×C的重疊patch，并通過兩層連續 3×3的CNN 層（步幅為2）構成的主干投影到C維嵌入空間中。

在各個階段之間的下采樣模塊由一個批歸一化的3×3的CNN 層（步幅為2）組成，將圖像分辨率減半。

此外，第1和第2階段中的CNN模塊，采用了通用的殘差模塊結構，具體如下：

其中：Conv3×3 表示3×3卷積操作；BN表示批歸一化（Batch Normalization）；GELU 是激活函數，表示 Gaussian Error Linear Unit；z^ 是經過卷積、批歸一化和激活函數處理后的中間結果；最后，z是通過卷積和批歸一化后的結果與原始輸入相加，形成殘差連接。

這種結構有助于緩解深層網絡訓練中的梯度消失問題，并提高模型的訓練效率。

Mamba架構

Mamba是結構化狀態空間序列模型的擴展，能夠通過可學習的隱狀態 h(t)，將一維連續輸入x(t)轉換為y(t)。該過程的公式如下：

其中，矩陣A，B，C是模型的參數。

離散化：為了提高計算效率，以上公式中的連續參數A，B和C需要轉化為離散參數。具體而言，假設時間步長為Δ，可以應用零階保持規則來獲取離散參數：

這種離散化方法能夠提升計算效率，便于在實際應用中實現 Mamba 模型。

使用離散參數代入到原方程：

此外，對于一個大小為T的輸入序列，可以用帶有卷積核K的全局卷積，進一步簡化上式中的輸出，具體如下

選擇性：Mamba 進一步擴展了S4公式，引入了一種選擇機制，使得模型能夠進行依賴于輸入的序列處理。這種機制使得模型的參數B 、C和Δ可以根據輸入動態調整，從而濾除無關信息。

設輸入X是TxC矩陣，其中 T 為序列長度，C為嵌入維度，第3和第4階段的第n層輸出可以按如下方式計算：

其中，NormNorm和MixerMixer分別表示層歸一化和 token 混合模塊的選擇。

層架構

在不失一般性的情況下，層歸一化（Layer Normalization）被用于 NormNorm。給定N層，前 N/2層使用 MambaVision混合模塊，而剩余的N/2層使用自注意力機制。

MambaVision 混合模塊:重新設計了原始的Mamba混合模塊，使其更適合視覺任務。

如圖3所示，首先將因果卷積（causal convolution）替換為常規卷積，因為因果卷積將信息限制在一個方向上，這對視覺任務來說不僅沒必要，而且局限性還很大。

此外，添加了一個不包含SSM（狀態空間模型）的對稱分支，該分支由額外的卷積和SiLU激活函數組成，以補償由于SSM的順序約束而可能丟失的內容。

然后，將兩個分支的輸出拼接起來，并通過最終的線性層進行投影。這種組合確保了最終的特征表示，同時包含順序信息和空間信息，從而充分利用了兩個分支的優勢。

注意到，每個分支的輸出被投影到一個大小為C/2的嵌入空間（即原始嵌入維度的一半），以保持與原始模塊設計相似的參數量。

給定輸入Xin，MambaVision混合模塊的輸出Xout計算如下：

其中，Linear(Cin,Cout)(?)表示一個線性層，輸入和輸出的嵌入維度分別為Cin和Cout；Scan是選擇性掃描操作（selective scan）；σ是激活函數，這里使用的是Sigmoid線性單元（SiLU；Conv和Concat 分別表示1D卷積和拼接操作。

實驗結果

表1展示了ImageNet-1K分類結果。具體來說，與不同類別的模型進行了比較，包括基于卷積的模型、基于 Transformer的模型、卷積-Transformer混合模型以及基于Mamba的模型，并證明新模型在ImageNet Top-1準確率和圖像處理能力方面大幅超越了之前的工作。

例如，與流行的模型如ConvNeXt和Swin Transformers相比，MambaVision-B（84.2%）優于 ConvNeXt-B（83.8%）和 SwinB（83.5%），同時在圖像處理能力上也有顯著優勢。

在與基于 Mamba 的模型比較時也觀察到了類似的趨勢。具體來說，盡管MambaVision-B（84.2%的圖像處理能力顯著更高，但仍優于 VMamba-B（83.9%）。

與同等規模的模型相比，MambaVision 型變體的FLOPs遠低于它們。例如，MambaVision-B 的GFLOPs比 MaxViT-B 少了56%。

表2展示在MS COCO數據集上的目標檢測和實例分割結果。

具體來說，訓練了不同檢測尺寸的模型，以進一步驗證 MambaVision 不同場景下的有效性。

通過簡單的Mask-RCNN檢測頭,預訓練的MambaVision-T骨干網絡，超過了 ConvNeXt-T和 Swin-T模型。

使用Cascade Mask-RCNN網絡時，MambaVision-T、MambaVision-S和MambaVision-B都超過了競爭對手。

表3展示了在ADE20K數據集上的語義分割基準測試。

對于這些實驗，使用了 UPerNet，以便與其他模型進行比較。

觀察到，MambaVision 模型在不同變體下超越了同等規模的競爭模型。

例如，MambaVision-T、MambaVision-S 和 MambaVision-B分別在mIoU上超越了Swin-T、Swin-S和Swin-B，提升幅度為+0.6、+0.6和+1.0。

盡管沒有對下游任務進行大量的超參數調優優化，這些結果仍然證明了MambaVision作為一種有前景的視覺任務骨干網絡的可行性，特別是在高分辨率設置下。

消融實驗和更多細節請參考原文。