新智元報道
新智元報道
編輯：KingHZ
【新智元導讀】CVPR 2025，混合新架構MambaVision來了！Mamba+Transformer混合架構專門為CV應用設計。MambaVision 在Top-1精度和影像吞吐量方面實現了新的SOTA，顯著超越了基於Transformer和Mamba的模型。
正如標題所言「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作為一種有前景的視覺任務骨幹網路的可行性，特別是在高解析度設定下。
消融實驗和更多細節請參考原文。
參考資料：
https://arxiv.org/abs/2407.08083
https://x.com/ahatamiz1/status/1894838302450295183