一種新的通用視覺(jué)主干模型Vision Mamba

Vision Mamba

最具潛力的下一代通用視覺(jué)主干網(wǎng)絡(luò)

CNN和Transformer常作為深度學(xué)習(xí)模型的首選基礎(chǔ)模塊，被應(yīng)用于各種場(chǎng)景，如文本、視覺(jué)、語(yǔ)音信號(hào)處理及其各種下游應(yīng)用。然而這兩個(gè)基礎(chǔ)模塊都有著其固有而互補(bǔ)的缺陷：CNN具有固定大小窗口的卷積核，使其計(jì)算量為線性，但也由此而面臨著窗口化局部感受野的缺陷，使其在全局場(chǎng)景感知和場(chǎng)景語(yǔ)義理解上之力;Transformer通過(guò)全局的注意力計(jì)算，使其具有長(zhǎng)上下文的全局感知能力，但其二次方復(fù)雜度的計(jì)算量使得在算力有限的端側(cè)設(shè)備上運(yùn)行變得很困難。

針對(duì)這一問(wèn)題，我們提出了VisionMamba，來(lái)打破線性復(fù)雜度與全局感受野不可兼得的困境?；谧匀徽Z(yǔ)言處理中的Mamba狀態(tài)空間模型SSM，我們?cè)O(shè)計(jì)了雙向SSM，并引I入了位置編碼來(lái)專(zhuān)門(mén)處理具有二維多向性的視覺(jué)信號(hào)。在各種分類(lèi)、檢測(cè)、分割任務(wù)中，Vim相比現(xiàn)有的視覺(jué)Transformer在精度上具有顯著提升，同時(shí)在計(jì)算和內(nèi)存效率上也有顯著改進(jìn)。例如，在進(jìn)行分辨率為1248x1248的批量推理時(shí)，Vim比DeiT快2.8倍，GPU內(nèi)存節(jié)省86.8%

? Vision Mamba 論文鏈接：

https://arxiv.org/abs/2401.09417

? 項(xiàng)目主頁(yè)：

https://github.com/hustvl/Vim

簡(jiǎn)介

本文的工作Vision Mamba[1]發(fā)表在ICML 2024。研究的問(wèn)題是如何設(shè)計(jì)新型神經(jīng)網(wǎng)絡(luò)來(lái)實(shí)現(xiàn)高效的視覺(jué)表示學(xué)習(xí)。該任務(wù)要求神經(jīng)網(wǎng)絡(luò)模型能夠在處理高分辨率圖像時(shí)既保持高性能，又具備計(jì)算和內(nèi)存的高效性。先前的方法主要依賴(lài)自注意力機(jī)制來(lái)進(jìn)行視覺(jué)表示學(xué)習(xí)，但這種方法在處理長(zhǎng)序列時(shí)速度和內(nèi)存使用上存在挑戰(zhàn)。論文提出了一種新的通用視覺(jué)主干模型Vision Mamba，簡(jiǎn)稱(chēng)Vim1，該模型使用雙向狀態(tài)空間模型（SSM）對(duì)圖像序列進(jìn)行位置嵌入，并利用雙向SSM壓縮視覺(jué)表示。在ImageNet[2]分類(lèi)、COCO[2]目標(biāo)檢測(cè)和ADE20k[3]語(yǔ)義分割任務(wù)中，Vim相比現(xiàn)有的視覺(jué)Transformer[4]（如DeiT[5]）在性能上有顯著提升，同時(shí)在計(jì)算和內(nèi)存效率上也有顯著改進(jìn)。例如，在進(jìn)行分辨率為1248×1248的批量推理時(shí)，Vim比DeiT快2.8倍，GPU內(nèi)存節(jié)省86.8%。這些結(jié)果表明，Vim能夠克服在高分辨率圖像理解中執(zhí)行Transformer樣式的計(jì)算和內(nèi)存限制，具有成為下一代視覺(jué)基礎(chǔ)模型主干的潛力。

圖1 本文所提出的Vision Mamba (Vim)和基于Transformer的DeiT模型進(jìn)行精度與效率對(duì)比：Vim在圖像分類(lèi)、目標(biāo)檢測(cè)、語(yǔ)義分割、實(shí)例分割任務(wù)上獲得了更好的精度，且在高清分辨率圖像處理上呈現(xiàn)出巨大的優(yōu)勢(shì)。

研究背景

圖像表示學(xué)習(xí)是計(jì)算機(jī)視覺(jué)領(lǐng)域的重要研究課題，其目的是通過(guò)模型學(xué)習(xí)從圖像中提取有意義的特征，從而應(yīng)用于各種視覺(jué)任務(wù)中。目前，視覺(jué)Transformer（Vision Transformer, ViT[4]）和卷積神經(jīng)網(wǎng)絡(luò)（Convolutional Neural Networks, CNNs）是圖像表示學(xué)習(xí)中最常用的方法。然而，這些方法在理論上存在一些局限性。

視覺(jué)Transformer利用自注意力機(jī)制能夠取得全局的感受野，在大規(guī)模自監(jiān)督預(yù)訓(xùn)練和下游任務(wù)中表現(xiàn)出色，但其自注意力機(jī)制在處理長(zhǎng)序列依賴(lài)和高分辨率圖像時(shí)，帶來(lái)了計(jì)算和內(nèi)存的巨大開(kāi)銷(xiāo)。具體而言，自注意力機(jī)制的計(jì)算復(fù)雜度是輸入的圖像塊序列長(zhǎng)度的平方，這使得其在處理高分辨率圖像時(shí)非常耗時(shí)且占用大量?jī)?nèi)存。盡管一些研究提出了改進(jìn)方法，如窗口注意力機(jī)制[6,7]，但這些方法雖然降低了復(fù)雜度，但導(dǎo)致感受野被局限在局部的窗口內(nèi)部，失去了原本全局感受野的優(yōu)勢(shì)。

另一方面，卷積神經(jīng)網(wǎng)絡(luò)在處理圖像時(shí)，通過(guò)使用固定大小的卷積核來(lái)提取局部特征。然而，卷積神經(jīng)網(wǎng)絡(luò)在捕捉全局上下文信息方面存在局限性，因?yàn)榫矸e核的感受野是有限的，雖然一些研究引入了金字塔結(jié)構(gòu)或大卷積核來(lái)增強(qiáng)全局信息提取能力，但這些改進(jìn)仍然無(wú)法完全克服CNN在處理長(zhǎng)序列依賴(lài)方面的不足。

在自然語(yǔ)言處理領(lǐng)域，Mamba[11]方法的出現(xiàn)給高效率長(zhǎng)序列建模帶來(lái)了很好的發(fā)展契機(jī)。Mamba是狀態(tài)空間模型（state space model, SSM）方法的最新演進(jìn)。Mamba提出了一種輸入自適應(yīng)的狀態(tài)空間模型，能夠更高質(zhì)量地完成序列建模任務(wù)。與此同時(shí)，該方法在處理長(zhǎng)序列建模問(wèn)題時(shí)有著次二次方的復(fù)雜度與更高的處理效率。然而，Mamba方法并不能夠直接應(yīng)用于視覺(jué)表征學(xué)習(xí)，因?yàn)镸amba方法是為自然語(yǔ)言領(lǐng)域的因果建模而設(shè)計(jì)的，它缺少對(duì)于二維空間位置的感知能力以及缺少全局的建模能力。

圖2 本文所提出的Vim模型的網(wǎng)絡(luò)構(gòu)架圖。

為了克服上述Transformer和CNN的理論局限性，啟發(fā)于自然語(yǔ)言處理領(lǐng)域Mamba的成功，本文提出了一種新的通用視覺(jué)主干模型——Vision Mamba (Vim)。該模型基于狀態(tài)空間模型[10]（State Space Models, SSMs），利用其在長(zhǎng)序列建模中的高效性，提供了一種新的視覺(jué)表示學(xué)習(xí)方法。該模型提出了雙向狀態(tài)空間模型來(lái)適配視覺(jué)特征的多方向性，并引入位置編碼來(lái)針對(duì)圖像單元進(jìn)行標(biāo)記。本文提出的Vim模型通過(guò)雙向SSM對(duì)圖像序列進(jìn)行位置嵌入和壓縮，不僅在ImageNet分類(lèi)任務(wù)上表現(xiàn)出色，還在COCO目標(biāo)檢測(cè)和ADE20k語(yǔ)義分割任務(wù)中展示了優(yōu)異的性能。與現(xiàn)有的視覺(jué)Transformer如DeiT相比，Vim在計(jì)算和內(nèi)存效率上有顯著提升。

Vision Mamba方法介紹

“

狀態(tài)空間模型

狀態(tài)空間模型，比如結(jié)構(gòu)化狀態(tài)空間序列模型[10]（S4）和Mamba[11]是啟發(fā)于連續(xù)系統(tǒng)，該系統(tǒng)通過(guò)隱藏狀態(tài)將一維函數(shù)或序列映射到。該系統(tǒng)使用作為演化參數(shù)，并使用和作為投影參數(shù)。連續(xù)系統(tǒng)的工作方式如下： ?

S4和Mamba是連續(xù)系統(tǒng)的離散版本，它們包含一個(gè)時(shí)間尺度參數(shù)，用于將連續(xù)參數(shù)A和B轉(zhuǎn)換為離散參數(shù)

和

。常用的方式是零階保持，其定義如下：

將

和

離散化后，使用步長(zhǎng)的離散版本可以重寫(xiě)為：

最后，模型可以使用全局的卷積來(lái)并行的計(jì)算：

其中

是輸入序列

的長(zhǎng)度，是結(jié)構(gòu)化的卷積核。 ?

“

Vision Mamba結(jié)構(gòu)

所提出的Vision Mamba如圖1所示。標(biāo)準(zhǔn)的Mamba模塊是為一維的文本序列所設(shè)計(jì)的。為了適配視覺(jué)信號(hào)，我們首先將二維圖像轉(zhuǎn)換為展平的二維圖像塊序列，其中是輸入圖像的尺寸，C是通道數(shù)，P是圖像塊的尺寸。接下來(lái)，我們將線性投影到大小為D的向量，并添加位置編碼

，如下所示：

其中

是

中的第

個(gè)圖像塊，

是可學(xué)習(xí)的投影變換矩陣。受ViT[4]的啟發(fā)，我們也使用類(lèi)別標(biāo)記來(lái)表示整個(gè)圖像塊序列。然后，我們將標(biāo)記序列輸入到Vim編碼器的第層，并得到輸出。最后我們對(duì)輸出類(lèi)別標(biāo)記進(jìn)行歸一化，并將其送入多層感知機(jī)（MLP）分類(lèi)頭以獲得最終類(lèi)別預(yù)測(cè)：

其中Vim是提出的視覺(jué)Mamba模塊，

是層數(shù)，

是歸一化層。

算法1：Vim模塊流程

輸入：圖像塊序列

輸出：圖像塊序列

“

Vim模塊

原始的Mamba模塊是為一維序列設(shè)計(jì)的，不適用于需要空間感知理解的視覺(jué)任務(wù)。我們創(chuàng)新性提出Vision Mamba編碼的基本構(gòu)建模塊Vim模塊，如圖2右側(cè)所示。具體來(lái)說(shuō)，像我們?cè)谒惴?中所展示的操作。輸入的標(biāo)記序列首先通過(guò)歸一化層進(jìn)行歸一化。接下來(lái)，我們將歸一化后的序列線性投影到維度大小為的和。然后，我們從前向和后向兩個(gè)方向處理。對(duì)于每個(gè)方向，我們首先對(duì)進(jìn)行一維卷積，得到。然后，我們將線性投影到。然后用于分別離散化得到和。最后我們通過(guò)SSM計(jì)算前向輸出和反向輸出，并通過(guò)進(jìn)行門(mén)控，并加在一起得到輸出標(biāo)記序列。 ?

“

效率優(yōu)化

Vim通過(guò)借助于Mamba的硬件友好的實(shí)現(xiàn)方式確保運(yùn)行的效率。優(yōu)化的關(guān)鍵思想是避免GPU的I/O瓶頸和內(nèi)存瓶頸。

IO高效性。高帶寬存儲(chǔ)器（HBM）和SRAM是GPU的兩個(gè)重要組成部分。其中，SRAM具有更大的帶寬，而HBM具有更大的存儲(chǔ)容量。標(biāo)準(zhǔn)的Vim的SSM操作在HBM上需要的I/O數(shù)量是O(BMEN)，其中B為批量大小，M為圖像塊序列長(zhǎng)度，E 表示擴(kuò)展?fàn)顟B(tài)維度，N 表示 SSM 維度。受到Mamba的啟發(fā)，Vim首先將O(BME+EN)字節(jié)的內(nèi)存從較慢的HBM讀取到較快的SRAM中。然后Vim在SRAM中獲取對(duì)應(yīng)的參數(shù)，并執(zhí)行SSM操作，最終將輸出結(jié)果寫(xiě)回HBM。此方法可以講I/O數(shù)量從O(BMEN)降低到O(BME+EN)從而大幅度提升效率。

內(nèi)存高效性。為了避免內(nèi)存不足問(wèn)題并在處理長(zhǎng)序列時(shí)降低內(nèi)存使用，Vim選擇了與 Mamba 相同的重計(jì)算方法。對(duì)于尺寸為 (B.M,E,N)的中間狀態(tài)來(lái)計(jì)算梯度，Vim在網(wǎng)絡(luò)的反向傳遞中重新計(jì)算它們。對(duì)于激活函數(shù)和卷積的中間激活值，Vim 也重新計(jì)算它們，以?xún)?yōu)化 GPU 的內(nèi)存需求，因?yàn)榧せ钪嫡加昧舜罅績(jī)?nèi)存，但重新計(jì)算速度很快。

計(jì)算高效性。Vim模塊中的SSM算法和Transformer中的自注意力機(jī)制都在自適應(yīng)地提供全局上下文方面起到了關(guān)鍵作用。給定一個(gè)視覺(jué)序列和默認(rèn)的設(shè)置。全局注意力機(jī)制和SSM的計(jì)算復(fù)雜度分別為：

其中，自注意力機(jī)制的計(jì)算復(fù)雜度和序列長(zhǎng)度成平方關(guān)系，而SSM的計(jì)算復(fù)雜度和序列長(zhǎng)度呈線性關(guān)系。這種計(jì)算效率使得Vim在處理具有長(zhǎng)序列長(zhǎng)度的千兆像素級(jí)別應(yīng)用時(shí)具有良好的擴(kuò)展性。

實(shí)驗(yàn)結(jié)果

該方法在標(biāo)準(zhǔn)的大型圖片分類(lèi)數(shù)據(jù)集ImageNet-1K上進(jìn)行驗(yàn)證。并將分類(lèi)訓(xùn)練好的模型作為預(yù)加載權(quán)重用于下游圖片密集型預(yù)測(cè)任務(wù)中去，如COCO數(shù)據(jù)集上的目標(biāo)檢測(cè)和實(shí)力分割任務(wù)， ADE20K上的像素級(jí)別的語(yǔ)義分割任務(wù)。

“

分類(lèi)對(duì)比

如表1與當(dāng)前主流的分類(lèi)模型對(duì)比Vim顯示出了相當(dāng)?shù)木?，將Vim和基于CNN、Transformer和SSM的主干網(wǎng)絡(luò)進(jìn)行比較，Vim顯示了相當(dāng)甚至更優(yōu)的性能。例如，在參數(shù)量相同的情況下Vim-Small的準(zhǔn)確率80.3%，比ResNet50[12]高出了4.1個(gè)百分點(diǎn)。與傳統(tǒng)的基于自注意力機(jī)制的ViT[4]相比，Vim在參數(shù)數(shù)量和準(zhǔn)確率上均有顯著提升。與視覺(jué)Transformer ViT高度優(yōu)化的變種DeiT相比，Vim在不同模型尺度上均以相似的參數(shù)數(shù)量取得了更好的精度。

如圖1所示，Vim的優(yōu)越的效率足以支持更細(xì)粒度的微調(diào)，在通過(guò)細(xì)粒度微調(diào)后，與基于SSM的S4ND-ViT-B[13]相比，Vim在參數(shù)數(shù)量減小3倍的情況下達(dá)到了相似的精度，Vim-Ti+，Vim-S+和Vim-B+的結(jié)果均有所提高。其中，Vim-S+甚至達(dá)到了與DeiT-B相似的效果。

表1ImageNet-1K分類(lèi)骨干網(wǎng)絡(luò)對(duì)比

“

語(yǔ)義分割對(duì)比

在ADE20K語(yǔ)義分割數(shù)據(jù)集上，我們將ImageNet-1K上訓(xùn)練好的權(quán)重加載到UperNet[14]分割器中，使用Vim作為骨干網(wǎng)絡(luò)進(jìn)行特征提取，如表3所示，Vim取得了相比于CNN網(wǎng)絡(luò)ResNet更少的參數(shù)量以及更高的精度，去Transformer模型DeiT相比，Vim取得了更優(yōu)的精度。

表2ADE20k語(yǔ)義分割對(duì)比

表3 COCO目標(biāo)檢測(cè)和實(shí)例分割對(duì)比

“

目標(biāo)檢測(cè)與實(shí)例分割對(duì)比

在COCO目標(biāo)檢測(cè)與實(shí)例分割數(shù)據(jù)集上，我們將ImageNet-1K上訓(xùn)練好的權(quán)重加載到Cascade-RCNN框架中，使用Vim作為骨干網(wǎng)絡(luò)進(jìn)行特征提取，如表3所示，Vim取得相對(duì)于Transformer的DeiT更好的檢測(cè)框精度和實(shí)例分割精度。值得注意的是，在高清圖像輸入的目標(biāo)檢測(cè)任務(wù)上，圖像輸入分辨率為1024×1024，由于Transformer的平方復(fù)雜度，需要將自注意力機(jī)制限制在固定大小的窗口內(nèi)， 而Vim得意于其線性復(fù)雜度，無(wú)需窗口化，可以進(jìn)行全局的視覺(jué)特征感知，從而取得了相對(duì)于表3中窗口化DeiT更好的精度。

“

消融實(shí)驗(yàn)

雙向SSM。如表4所示，雙向SSM相較于原本的單向SSM取得了更高的分類(lèi)精度， 且在下游的密集型預(yù)測(cè)任務(wù)上取得更為顯著的優(yōu)勢(shì)。這一結(jié)果顯示了本文提出的雙向設(shè)計(jì)對(duì)于視覺(jué)特征學(xué)習(xí)的必要性與重要性。

表4 雙向SSM建模消融實(shí)驗(yàn)

分類(lèi)策略。在表5中，我們探索了以下幾種分類(lèi)策略：

表5 分類(lèi)策略消融實(shí)驗(yàn)

·Mean pool，將最后Vision Mamba編碼器輸出的特征進(jìn)行平均池化。

·Max pool，將最后Vision Mamba編碼器輸出的特征進(jìn)行最大化池化。

·Head class token，將類(lèi)別標(biāo)記詞元置于圖像塊序列頭部。

·Double class token，將類(lèi)別標(biāo)記詞元置于圖像塊序列兩端。

·Middle class token，將類(lèi)別標(biāo)記詞元置于圖像塊序列中間。

如表5所示，實(shí)驗(yàn)結(jié)果表明，中間類(lèi)別標(biāo)記策略能夠充分利用SSM的循環(huán)特性和ImageNet中的中心對(duì)象先驗(yàn)，展示了最佳的top-1準(zhǔn)確率76.1。

總結(jié)

該論文提出了Vision Mamba (Vim)，以探索最新的高效狀態(tài)空間模型Mamba作為通用視覺(jué)主干網(wǎng)絡(luò)。與以往用于視覺(jué)任務(wù)的狀態(tài)空間模型采用混合架構(gòu)或等效的全局二維卷積核不同，Vim以序列建模的方式學(xué)習(xí)視覺(jué)表示，并未引入圖像特定的歸納偏置。得益于所提出的雙向狀態(tài)空間建模，Vim實(shí)現(xiàn)了數(shù)據(jù)依賴(lài)的全局視覺(jué)上下文，并具備與Transformer相同的建模能力，同時(shí)計(jì)算復(fù)雜度更低。受益于Mamba的硬件感知設(shè)計(jì)，Vim在處理高分辨率圖像時(shí)的推理速度和內(nèi)存使用顯著優(yōu)于ViTs。在標(biāo)準(zhǔn)計(jì)算機(jī)視覺(jué)基準(zhǔn)上的實(shí)驗(yàn)結(jié)果驗(yàn)證了Vim的建模能力和高效性，表明Vim具有成為下一代視覺(jué)主干網(wǎng)絡(luò)的巨大潛力。