Vision Mamba：速度與內(nèi)存的雙重突破

Vision Mamba 不是個普通模型。

號稱「全面包圍 Transformer」的 Mamba，推出不到兩個月就有了高性能的視覺版。 來自華中科技大學(xué)、地平線、智源人工智能研究院等機構(gòu)的研究者提出了 Vision Mamba（Vim）。

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

接下來我們看看論文內(nèi)容。 Mamba 的提出帶動了研究者對狀態(tài)空間模型（state space model，SSM）興趣的增加，不同于 Transformer 中自注意力機制的計算量會隨著上下文長度的增加呈平方級增長，由于 SSM 擅長捕捉遠(yuǎn)程依賴關(guān)系，因而開始受到大家追捧。 在此期間，一些基于 SSM 的方法如線性狀態(tài)空間層（LSSL）、結(jié)構(gòu)化狀態(tài)空間序列模型（S4）、對角狀態(tài)空間（DSS）和 S4D 都被研究者提出來，用于處理各種序列數(shù)據(jù)，特別是在建模遠(yuǎn)程依賴關(guān)系方面。 Mamba 將時變參數(shù)納入 SSM 中，并提出了一種硬件感知算法來實現(xiàn)高效的訓(xùn)練和推理。Mamba 卓越的擴展性能表明它在語言建模方面是 Transformer 有前途的替代品。 然而，到目前為止，研究者還尚未在視覺任務(wù)中探索出通用的基于純 SSM 的骨干網(wǎng)絡(luò)。 受 Mamba 在語言建模方面成功的激勵，研究者開始設(shè)想能否將這種成功從語言轉(zhuǎn)移到視覺，即用先進的 SSM 方法設(shè)計通用且高效的視覺主干。

然而，由于 Mamba 特有的架構(gòu)，需要解決兩個挑戰(zhàn)，即單向建模和缺乏位置感知。 為了應(yīng)對這些問題，研究者提出了 Vision Mamba (Vim) 塊，它結(jié)合了用于數(shù)據(jù)依賴的全局視覺上下文建模的雙向 SSM 和用于位置感知視覺識別的位置嵌入。 與其他基于 SSM 的視覺任務(wù)模型相比，Vim 是一種基于純 SSM 的方法，并以序列方式對圖像進行建模。與基于 Transformer 的 DeiT 相比，Vim 在 ImageNet 分類上取得了優(yōu)越的性能。此外，Vim 在 GPU 內(nèi)存和高分辨率圖像的推理時間方面更加高效。 方法介紹 Vision Mamba (Vim) 的目標(biāo)是將先進的狀態(tài)空間模型 (SSM)，即 Mamba 引入到計算機視覺。 Vim 的概述如圖 2 所示，標(biāo)準(zhǔn)的 Mamba 是為 1-D 序列設(shè)計的。為了處理視覺任務(wù)，首先需要將二維圖像轉(zhuǎn)換成展開的 2-D patch?

。式中 (H, W) 為輸入圖像的大小，C 為通道數(shù)，P 為圖像 patch 的大小。接下來，需要將 x_p 線性投影到大小為 D 的向量上，并添加位置嵌入

得到如下公式：

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

架構(gòu)細(xì)節(jié) 架構(gòu)的超參數(shù)如下：

L：塊數(shù)

D：隱藏狀態(tài)維度

E：擴展?fàn)顟B(tài)維度

N：SSM 維度

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

實驗

該研究在 ImageNet-1K 數(shù)據(jù)集上對 Vim 進行了基準(zhǔn)測試。

圖像分類

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

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

語義分割 如表 2 所示，Vim 在不同尺度上始終優(yōu)于 DeiT：Vim-Ti 比 DeiT-Ti 高 1.0 mIoU，Vim-S 比 DeiT-S 高 0.9 mIoU。與 ResNet-101 主干網(wǎng)絡(luò)相比，Vim-S 以減少近 2 倍的參數(shù)實現(xiàn)了相同的分割性能。

為了進一步評估研究方法在下游任務(wù)上（即分割、檢測和實例分割）的效率，本文將骨干網(wǎng)與常用的特征金字塔網(wǎng)絡(luò)（FPN）模塊結(jié)合起來，并對其 FPS 和 GPU 內(nèi)存進行基準(zhǔn)測試。 如圖 3 和圖 4 所示，盡管該研究在主干網(wǎng)上附加了一個 heavy FPN，但效率曲線與純主干網(wǎng)（圖 1）的比較結(jié)果相似。

目標(biāo)檢測和實例分割 表 3 使用 Cascade Mask R-CNN 框架對 Vim-Ti 和 DeiT-Ti 進行了比較。Vim-Ti 超過 DeiT-Ti 1.3 box AP 和 1.1 mask AP。

審核編輯：黃飛