地平線(xiàn)Vision Mamba：超越ViT，最具潛力的下一代通用視覺(jué)主干網(wǎng)絡(luò )

? Vision Mamba 論文鏈接：

https://arxiv.org/abs/2401.09417

? 項目主頁(yè)：

https://github.com/hustvl/Vim

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

表3  COCO目標檢測和實(shí)例分割對比

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

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

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

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

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

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

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

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

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

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

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

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

參考文獻：

[1] Zhu L, Liao B, Zhang Q, et al. Vision mamba: Efficient visual representation learning with bidirectional state space model. In ICML 2024.

[2] Deng J, Dong W, Socher R, et al. Imagenet: A large-scale hierarchical image database. In CVPR 2009.[3] Lin T Y, Maire M, Belongie S, et al. Microsoft coco: Common objects in context. In ECCV 2014.

[4] Dosovitskiy A, Beyer L, Kolesnikov A, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR 2021.

[5] Touvron H, Cord M, Douze M, et al. Training data-efficient image transformers & distillation through attention. In ICML 2021.

[6] Liu Z, Lin Y, Cao Y, et al. Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV 2021.

[7] Li Y, Mao H, Girshick R, et al. Exploring plain vision transformer backbones for object detection. In ECCV 2022.

[8] Vaswani A, Shazeer N, Parmar N, et al. Attention is all you need. In NeurIPS 2017.

[9] LeCun Y, Bottou L, Bengio Y, et al. Gradient-based learning applied to document recognition. In Proceedings of the IEEE 1998.

[10] Gu A, Goel K, Ré C. Efficiently modeling long sequences with structured state spaces. In ICLR 2022.

[11] Gu A, Dao T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.

[12] He K, Zhang X, Ren S, et al. Deep residual learning for image recognition. In CVPR 2016.

[13] Nguyen E, Goel K, Gu A, et al. S4nd: Modeling images and videos as multidimensional signals with state spaces. In NeurIPS 2022.

[14] Xiao T, Liu Y, Zhou B, et al. Unified perceptual parsing for scene understanding. In ECCV 2018.

[15] Cai Z, Vasconcelos N. Cascade r-cnn: Delving into high quality object detection. In CVPR 2018.

[16] Liao B, Wang X, Zhu L, et al. ViG: Linear-complexity Visual Sequence Learning with Gated Linear Attention. arXiv preprint arXiv:2405.18425, 2024.

[17] Zhu L, Huang Z, Liao B, et al. DiG: Scalable and Efficient Diffusion Models with Gated Linear Attention. arXiv preprint arXiv:2405.18428, 2024.

[18] Peebles W, Xie S. Scalable diffusion models with transformers. In CVPR 2023.