視覺Mamba來了：速度提升2.8倍，內存能省87%

作者：機器之心 2024-01-19 13:41:00

來自華中科技大學、地平線、智源人工智能研究院等機構的研究者提出了 Vision Mamba（Vim）。

號稱「全面包圍 Transformer」的 Mamba，推出不到兩個月就有了高性能的視覺版。

本周四，來自華中科技大學、地平線、智源人工智能研究院等機構的研究者提出了 Vision Mamba（Vim）。

論文地址：https://arxiv.org/pdf/2401.09417.pdf
項目地址：https://github.com/hustvl/Vim
論文標題：Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

效果如何呢？在 ImageNet 分類任務、COCO 對象檢測任務和 ADE20k 語義分割任務上，與 DeiT 等成熟的視覺 Transformers 相比，Vim 實現了更高的性能，同時還顯著提高了計算和內存效率。例如，在對分辨率為 1248×1248 的圖像進行批量推理提取特征時，Vim 比 DeiT 快 2.8 倍，并節省 86.8% 的 GPU 內存。結果表明，Vim 能夠克服對高分辨率圖像執行 Transformer 式理解時的計算和內存限制，并且具有成為視覺基礎模型的下一代骨干的巨大潛力。

接下來我們看看論文內容。

Mamba 的提出帶動了研究者對狀態空間模型（state space model，SSM）興趣的增加，不同于 Transformer 中自注意力機制的計算量會隨著上下文長度的增加呈平方級增長，由于 SSM 擅長捕捉遠程依賴關系，因而開始受到大家追捧。

在此期間，一些基于 SSM 的方法如線性狀態空間層（LSSL）、結構化狀態空間序列模型（S4）、對角狀態空間（DSS）和 S4D 都被研究者提出來，用于處理各種序列數據，特別是在建模遠程依賴關系方面。

Mamba 將時變參數納入 SSM 中，并提出了一種硬件感知算法來實現高效的訓練和推理。Mamba 卓越的擴展性能表明它在語言建模方面是 Transformer 有前途的替代品。

然而，到目前為止，研究者還尚未在視覺任務中探索出通用的基于純 SSM 的骨干網絡。

受 Mamba 在語言建模方面成功的激勵，研究者開始設想能否將這種成功從語言轉移到視覺，即用先進的 SSM 方法設計通用且高效的視覺主干。然而，由于 Mamba 特有的架構，需要解決兩個挑戰，即單向建模和缺乏位置感知。

為了應對這些問題，研究者提出了 Vision Mamba (Vim) 塊，它結合了用于數據依賴的全局視覺上下文建模的雙向 SSM 和用于位置感知視覺識別的位置嵌入。

與其他基于 SSM 的視覺任務模型相比，Vim 是一種基于純 SSM 的方法，并以序列方式對圖像進行建模。與基于 Transformer 的 DeiT 相比，Vim 在 ImageNet 分類上取得了優越的性能。此外，Vim 在 GPU 內存和高分辨率圖像的推理時間方面更加高效。

方法介紹

Vision Mamba (Vim) 的目標是將先進的狀態空間模型 (SSM)，即 Mamba 引入到計算機視覺。

Vim 的概述如圖 2 所示，標準的 Mamba 是為 1-D 序列設計的。為了處理視覺任務，首先需要將二維圖像轉換成展開的 2-D patch 。式中 (H, W) 為輸入圖像的大小，C 為通道數，P 為圖像 patch 的大小。接下來，需要將 x_p 線性投影到大小為 D 的向量上，并添加位置嵌入得到如下公式：

Vim 塊

原始的 Mamba 塊是為一維序列設計的，不適合需要空間感知理解的視覺任務。Vim 塊集成了用于視覺任務的雙向序列建模，Vim 塊如上圖 2 所示。

Vim 塊的操作算法如下所示。

架構細節

架構的超參數如下：

L：塊數
D：隱藏狀態維度
E：擴展狀態維度
N：SSM 維度

遵循 ViT 和 DeiT，該研究首先采用 16×16 內核大小的投影層來獲得非重疊 patch 嵌入的一維序列。隨后直接堆疊 L 個 Vim 塊。默認情況下塊數 L 設置為 24，SSM 維度 N 設置為 16。為了與 DeiT 系列模型大小保持一致，該研究將小（ tiny）尺寸變體的隱藏狀態維度 D 設置為 192，將擴展狀態維度 E 設置為 384。對于小（small）尺寸變體，該研究將 D 設置為 384，將 E 設置為 768。

實驗

該研究在 ImageNet-1K 數據集上對 Vim 進行了基準測試。

圖像分類

表 1 將 Vim 與基于 ConvNet、基于 Transformer 和基于 SSM 的骨干網絡進行了比較。與基于 ConvNet 的 ResNet 相比，Vim 表現出更優越的性能。例如，當參數大致相似時，Vim-Small 的 top-1 準確率達到 80.3，比 ResNet50 高 4.1 個百分點。與傳統的基于自注意力的 ViT 相比，Vim 在參數數量和分類準確率方面都有相當大的優勢。與高度優化的 ViT 變體（即 DeiT ）相比，VimTiny 比 DeiT-Tiny 高 0.9 個點，Vim-Small 比 DeiT 高 0.5 個點。與基于 SSM 的 S4ND-ViTB 相比，Vim 以減少 3 倍的參數實現了類似的 top-1 準確率。

圖 1 (b) 和 (c) 比較了小型 Vim 和 DeiT 的 FPS 和 GPU 內存。隨著圖像分辨率的提高，Vim 在速度和內存方面表現出更好的效率。具體來說，當圖像大小為 512 時，Vim 實現了與 DeiT 相似的 FPS 和內存。當圖像大小增長到 1248 時，Vim 比 DeiT 快 2.8 倍，并節省 86.8% 的 GPU 內存。Vim 在序列長度上的線性擴展的顯著優勢使其為高分辨率下游視覺應用和長序列多模態應用做好了準備。