視覺Mamba來了：速度提升2.8倍，內(nèi)存能省87%

作者：機(jī)器之心 2024-01-19 13:41:00

來自華中科技大學(xué)、地平線、智源人工智能研究院等機(jī)構(gòu)的研究者提出了 Vision Mamba（Vim）。

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

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

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

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

接下來我們看看論文內(nèi)容。

Mamba 的提出帶動(dòng)了研究者對(duì)狀態(tài)空間模型（state space model，SSM）興趣的增加，不同于 Transformer 中自注意力機(jī)制的計(jì)算量會(huì)隨著上下文長度的增加呈平方級(jí)增長，由于 SSM 擅長捕捉遠(yuǎn)程依賴關(guān)系，因而開始受到大家追捧。

在此期間，一些基于 SSM 的方法如線性狀態(tài)空間層（LSSL）、結(jié)構(gòu)化狀態(tài)空間序列模型（S4）、對(duì)角狀態(tài)空間（DSS）和 S4D 都被研究者提出來，用于處理各種序列數(shù)據(jù)，特別是在建模遠(yuǎn)程依賴關(guān)系方面。

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

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

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

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

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

方法介紹

Vision Mamba (Vim) 的目標(biāo)是將先進(jìn)的狀態(tài)空間模型 (SSM)，即 Mamba 引入到計(jì)算機(jī)視覺。

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

Vim 塊

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

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

架構(gòu)細(xì)節(jié)

架構(gòu)的超參數(shù)如下：

L：塊數(shù)
D：隱藏狀態(tài)維度
E：擴(kuò)展?fàn)顟B(tài)維度
N：SSM 維度

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

實(shí)驗(yàn)

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

圖像分類

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

圖 1 (b) 和 (c) 比較了小型 Vim 和 DeiT 的 FPS 和 GPU 內(nèi)存。隨著圖像分辨率的提高，Vim 在速度和內(nèi)存方面表現(xiàn)出更好的效率。具體來說，當(dāng)圖像大小為 512 時(shí)，Vim 實(shí)現(xiàn)了與 DeiT 相似的 FPS 和內(nèi)存。當(dāng)圖像大小增長到 1248 時(shí)，Vim 比 DeiT 快 2.8 倍，并節(jié)省 86.8% 的 GPU 內(nèi)存。Vim 在序列長度上的線性擴(kuò)展的顯著優(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ù)實(shí)現(xiàn)了相同的分割性能。