VMamba: Visual State Space Model

摘要

卷积神经网络(CNNs)和Vision Transformers (ViTs)是视觉表示学习最流行的两种基础模型。CNNs在线性复杂度和图像分辨率方面表现出显著的可扩展性，而ViTs在拟合能力方面超过了CNNs，尽管它要应对二次复杂度。进一步的观察表明，ViTs通过整合全局感受野和动态权重实现了卓越的视觉建模性能。这种观察促使我们提出一种新的体系结构，在继承这些组件的同时提高计算效率。为此，我们从最近引入的状态空间模型中获得灵感，提出了视觉状态空间模型(VMamba)，该模型在不牺牲全局感受野的情况下实现了线性复杂度。为了解决方向敏感问题，我们引入了交叉扫描模块(Cross-Scan Module, CSM)来遍历空间域，并将任何非因果视觉图像转换为有序的patch序列。大量的实验结果证实VMamba不仅在各种视觉感知任务中展示了有前途的能力，而且随着图像分辨率的提高，VMamba比既定的基准显示出更明显的优势。源代码可以从https://github.com/MzeroMiko/VMamba获得。

1介绍

视觉表示学习是计算机视觉中最基础的研究课题之一，自深度学习时代开始以来，视觉表示学习取得了重大突破。卷积神经网络(Convolution Neural Networks, CNNs)[38,19,22,29,42]和视觉变形器(Vision transformer, ViTs)[10,28,45,56]这两类主要的深度基础模型已被广泛应用于各种视觉任务中。虽然两者在计算表达性视觉表示方面都取得了显著的成功，但ViTs通常比CNNs表现出更好的性能，这可以归因于注意力机制促进的全局感受野和动态权重。

然而，在图像大小方面，注意力机制需要二次复杂度，因此在处理下游密集的预测任务时，如目标检测、语义分割等，会导致昂贵的计算开销。为了解决这个问题，大量的工作已经致力于通过限制计算窗口[43]的大小或步长来提高注意的效率，尽管代价是限制感受野的规模。这促使我们设计一个具有线性复杂度的新的视觉基础模型，同时仍然保留与全局感受野和动态权重相关的优势。

从最近提出的状态空间模型中汲取灵感[12,34,47]，我们引入了视觉状态空间模型(被称为VMamba)来进行有效的视觉表示学习。VMamba在有效降低注意力复杂性方面的成功背后的关键概念继承自选择性扫描空间状态顺序模型(S6)[12]，该模型最初设计用于处理自然语言处理(NLP)任务。与传统的注意计算方法不同，S6使一维数组(如文本序列)中的每个元素都可以通过压缩的隐藏状态与之前扫描的任意样本进行交互，有效地将二次复杂度降为线性。

然而，由于视觉数据的非因果性质，直接将这种策略应用于一个补丁化（patchified）、平面化的图像将不可避免地导致感受野受限，因为无法估计与未扫描补丁的关系。我们将这个问题称为“方向敏感”问题，并建议通过新引入的跨扫描模块(CSM)来解决这个问题。CSM并不是单向地(纵向或纵向)遍历图像特征图的空间域，而是采用四向扫描策略，即从整个特征图的四个角向相反位置扫描(如图2 (b)所示)。该策略保证了特征图中的每个元素在不同方向上整合了所有其他位置的信息，从而在不增加线性计算复杂度的情况下呈现出全局的感受野。

图2:信息流的比较:注意力机制vs.跨扫描模块(CSM)(a)注意力机制将所有像素统一整合为中心像素，导致复杂度O(N2)。(b) CSM集成了左上、右下、右上、左下的像素，复杂度为O(N)。

在不同的视觉任务上进行了大量的实验来验证VMamba的有效性。如图1所示，与包括Resnet[19]、ViT[10]和Swin[28] 1在内的基准视觉模型相比，VMamba模型在ImageNet-1K上显示出更好的或至少具有竞争力的性能。我们还报告了下游密集预测任务的结果。例如VMambaTiny/Small/Base(分别具有22/44/75 M参数)使用MaskRCNN检测器(1×训练计划)在COCO上实现了46.5%/48.2%/48.5%的mAP，使用512 × 512输入的UperNet在ADE20K上实现了47.3%/49.5%/50.0%的mIoU，展现了其作为一个强大的基础模型的潜力。此外，当使用更大的图像作为输入时，ViT的FLOPs比CNN模型的FLOPs增加得更快，尽管通常仍然表现出更好的性能。然而，有趣的是VMamba本质上是一个基于Transformer架构的基础模型，它能够在FLOPs稳步增加的情况下获得与ViT相当的性能。

图1:在ImageNet-1K上的性能比较，VMamba系列与流行的同行相比，达到了卓越的top-1精度。我们注意到所提出的VMamba能够显示全局有效接收域(ERF)，具有线性复杂度的动态权值。

我们总结这些贡献如下:

•我们提出了VMamba，一个具有全局感受野和动态权值的视觉状态空间模型，用于视觉表示学习。VMamba为视觉基础模型提供了一种新的选择，它超越了现有的CNNs和ViTs选择。

•引入了交叉扫描模块(CSM)，弥补了一维阵列扫描和二维平面遍历之间的差距，在不影响感受野的情况下，促进了S6向视觉数据的扩展。

•在没有弄虚作假的情况下，我们表明VMamba在各种视觉任务中，包括图像分类、对象检测和语义分割，都能取得很好的结果。这些发现强调了VMamba作为一个鲁棒的视觉基础模型的潜力。

2相关工作

深度神经网络极大地推进了机器视觉感知的研究。目前流行的视觉基础模型主要有两种类型，即CNNs[23, 38, 41, 19, 42]和ViTs[10, 28, 48, 9, 6, 56]。近年来，状态空间模型(SSMs)[12,34, 47]的成功证明了其在高效长序列建模中的有效性，引起了NLP和CV群体的广泛关注。我们的研究坚持了这一工作方向，并提出了VMamba，这是一种基于ssm的用于视觉领域数据建模的体系结构。VMamba与CNNs和ViTs一起，为社区提供了另一种基础模型。

卷积神经网络(CNNs)

卷积神经网络(CNNs)是视觉感知史上具有里程碑意义的模型。早期基于CNNs的模型[25,23]是为基本任务而设计的，如识别手写数字[24]和对字符类别[55]进行分类。卷积核将CNNs的特点封装在卷积核中，利用感受野从图像中获取感兴趣的视觉信息。在强大的计算设备(GPU)和大规模数据集[7]的帮助下，人们提出了越来越深入的[38,41,19,22]和高效的模型[20,42,52,36]，以提高各种视觉任务的性能。除了这些努力之外，人们还提出了更高级的卷积算子[4,21,53,5]或更高效的网络架构[59,3,51,20]。

Vision Transformers (ViTs)

Vision transformer (ViTs)改编自NLP社区，展示了一个强大的视觉任务感知模型，并迅速发展为最有前途的视觉基础模型之一。早期的基于虚拟现实的模型通常需要大规模的数据测试[10]，并以普通配置出现[54,58,1,31]。后来，DeiT[45]采用训练技术来解决优化过程中遇到的挑战，后续的研究倾向于将视觉感知的归纳偏差纳入网络设计。例如，社区提出分层的ViTs[28, 9, 48, 31, 56, 44, 6, 8, 57]来逐渐降低整个主干的特征分辨率。此外，其他研究也提出利用CNNs的优势，如引入卷积运算[49,6,46]，结合CNN和ViT模块设计混合架构[6,40,31]等。

状态空间模型(SSMs)

状态空间模型(SSMs)是最近提出的模型，将状态空间转换引入深度学习[16,15,39]。受控制系统中连续状态空间模型的启发，结合HiPPO[13]初始化，LSSL[16]展示了处理长期依赖问题的潜力。然而，由于状态表示带来的限制计算和内存需求，LSSL在实践中是不可行的。为了解决这一问题，S4[15]提出将参数归一化为对角结构。从那以后，出现了许多具有不同结构的结构化状态空间模型，如复杂对角结构[17,14]、支持[39]的多输入多输出、对角分解加低秩操作[18]、选择机制[12]。然后将这些模型集成到大型表示模型中[34,33,11]。

这些模型主要关注状态空间模型如何应用于语言和语音等长期和非常规数据，如语言理解[33,34]、基于内容的推理[12]、像素级1-D图像分类[15]，很少有人关注视觉识别。与我们最相似的作品是S4ND[35]。S4ND是第一项将状态空间机制应用于可视化任务的研究，显示了其与ViT[10]在性能上的竞争潜力。但S4ND对S4模型进行了简单的扩展，无法有效地获取依赖于输入的图像信息。我们证明，通过mamba[12]引入的选择性扫描机制，VMamba能够匹配现有流行的视觉基础模型，如ResNet[19]、ViT[10]、swin[27]和convnext[29]，这表明VMamba有成为基础模型的强大潜力。

3方法

在本节中，我们首先介绍VMamba相关的初步概念，包括状态空间模型、离散过程和选择性扫描机制。然后，我们提供了作为VMamba核心元素的2D状态空间模型的详细规范。最后，我们对VMamba整体架构进行了全面的讨论。

3.1初步准备

状态空间模型

状态空间模型(SSMs)通常被认为是线性时不变系统，它将激励映射到响应。在数学上，这些模型通常表示为线性常微分方程(ode) (Eq. 1)，对于状态大小N,其中参数包括, ，跳跃连接。

离散化

状态空间模型(SSMs)作为连续时间模型，在集成到深度学习算法中面临着巨大的挑战。为了克服这个障碍，离散化过程变得非常必要。

离散化的主要目标是将ODE转换为离散函数。这种转换对于使模型与输入数据中包含的底层信号的采样率保持一致至关重要，从而实现计算效率高的操作[16]。考虑输入，利用零阶保持规则离散如下:一个长度为L的信号流中的采样向量，一个长度为L的信号流中的采样向量，利用零阶保持规则将ODE (Eq. 1)离散如下:

其中，，。在实践中，在[12]之后，我们使用一阶泰勒级数来细化的近似:

选择性扫描机制

与主要关注线性时不变(LTI) SSM的主流方法不同，VMamba通过将选择性扫描机制(S6)[12]作为核心SSM算子而与众不同。在S6中，矩阵, ，是从输入数据得到的。这意味着S6通过嵌入上下文信息作为输入，从而确保了该机制中权重的动态性。

3.2 2D选择性扫描

尽管S6有其独特的特点，但它会对输入的数据进行处理，因此只能捕获被扫描的部分数据中的信息。这很自然地将S6与涉及时间数据的NLP任务结合起来，但在适应非因果数据(如图像、图表、集合等)时面临重大挑战。这个问题的一个直接的解决方案是沿着两个不同的方向(即，向前和向后)扫描数据，允许它们相互补充感受野，而不增加计算复杂度。

尽管图像具有非因果性质，但它与文本的不同之处在于包含了二维空间信息(如局部纹理和全局结构)。为了解决这个问题，S4ND[35]建议用卷积重新构造SSM，通过向量积直接将卷积核从1-D扩展到2-D。然而，这样的修改会阻碍动态权重(即，输入相互独立)，从而导致基于上下文的数据建模能力的丧失。因此，我们选择通过选择性扫描方法[12]来保留动态权重，不幸的是，这不允许我们遵循[35]和对卷积操作进行积分操作。

为了解决这个问题，我们提出了跨扫描模块(CSM)，如图2所示。我们选择将图像块沿行和列展开成序列(扫描展开)，然后沿着四个不同的方向继续扫描:从左上到右下、从右下到左上、从右上到左下、从左下到右上。通过这种方式，任何像素(如图2中的中心像素)都集成了来自不同方向的所有其他像素的信息。然后，我们将每个序列重新塑造为单个图像，并将所有序列合并为一个新的图像，如图3所示(扫描合并)。

图3:图像的2D-选择扫描示意图。我们首先使用CSM(扫描扩展)扫描一个图像。然后通过S6块分别处理四个结果特征，并将四个输出特征合并(扫描合并)，以构建最终的2D特征图。

S6与CSM的集成，称为S6块，作为构建Visual State Space (VSS)块的核心元素，VSS块是构成vmamba的基本构建块(在下一小节中进一步详细介绍)。我们强调，S6区块继承了选择性扫描机制的线性复杂性，同时保留了一个全局的感受野，这与我们构建这样一个视觉模型的动机一致。

3.3 VMamba模型

3.3.1整体架构

VMamba-Tiny的架构概述如图4 (a)所示。VMamba在开始过程时，使用类似于ViTs的stem module将输入图像分割成小块，但没有进一步将小块扁平化成一维序列。这种修改保留了图像的二维结构，得到了尺寸为H/4 × W/4 ×的特征图。

VMamba将几个VSS块堆叠，处理特征图上，保持相同的维数，构成Stage 1。VMamba中的层次表示是通过patch合并操作[27]对“Stage 1”中的特征图进行下采样来构建的。随后，涉及到更多的VSS块，输出分辨率为H/8 ×W/8，形成“Stage 2”。重复此过程以创建“阶段3”和“阶段4”，其分辨率分别为H/16×W/16和H/32×W/32。所有这些阶段共同构建了类似于流行的CNN模型[19,22,41,29,42]和一些ViTs[27,48,6, 56]的层次表示。最终的体系结构可以作为实际应用程序中具有类似需求的其他视觉模型的通用替代。

我们以三种不同的尺度开发VMamba，即VMamba-Tiny, VMamba-Small, 和VMamba-
Base (分别称为VMamba-T, VMamba-S, and VMamba-B, respectively)。详细的体系结构规范见表1。使用输入大小为224×224的图像评估所有模型的FLOPs。其他的体系结构，如大规模模型，将在未来的更新中引入。

表1:VMamba系列的体系结构概述。下采样通过阶段1、2、3的patch merge操作进行。术语“线性”指的是线性层。“DWConv”表示逐深度卷积操作。所提出的2D-选择性扫描被标记为“SS2D”。

3.3.2 VSS块

VSS块结构如图4 (b)所示，输入经过初始线性嵌入层，输出分为两个信息流。在进入核心SS2D模块之前，一个流通过一个3 × 3的逐深度卷积层，然后是Silu激活函数[37]。SS2D的输出经过一层层归一化层，然后添加到其他信息流的输出中，这些信息流经历了Silu函数激活。这种组合产生了VSS块的最终输出。

不同于vision transformers，我们避免使用位置嵌入偏置在VMamba，因为它的因果性质。我们的设计与典型的vision transformers结构不同，我们采用如下操作顺序:Norm→attention→Norm→MLP在一个块中，并且丢弃了MLP操作。因此，VSS块比ViT块要浅，这允许我们以类似的总模型深度来堆叠更多的块。

4 实验

在本节中，我们将执行一系列实验来评估和比较VMamba和主流的模型，包括CNNs和vision transformers。我们的评估涵盖了不同的任务，包括ImageNet-1K上的图像分类，COCO上的对象检测，以及ADE20K上的语义分割。随后，我们进行了分析实验，对VMamba的体系结构有了更深入的了解。

4.1 ImageNet-1K图像分类

设置

我们评估了VMamba在ImageNet-1K[7]上的分类性能。在[27]中配置之后，VMamba-T/S/B从头开始训练300个epoch(前20个epoch进行预热学习)，使用1024个批处理大小。训练过程包含了AdamW优化器，beta设置为(0.9,0.999)，动量为0.9，一个余弦衰减学习速率调度程序，初始学习速率为，权值衰减为0.05。额外的技术，如标签平滑(0.1)和指数移动平均(EMA)也被使用。除此之外，没有进一步的训练技术。

结果

表2总结了ImageNet-1K上的结果，将VMamba与流行的CNN模型和视觉变形器进行了比较。比较表明，在FLOPs相似的情况下，VMamba-T实现了82.2%的性能，比RegNetY-4G高出2.2%，比DeiT-S高出2.4%，比Swin-T高出0.9%。值得注意的是，VMamba在小型和常规规模的模型中仍然保持较好的性能优势。例如，在小尺度上，VMamba-S获得了83.5%的top-1精度，比RegNetY-8G高出1.8%，比Swin-S高出0.5%。与此同时，VMamba-B达到83.2%的top-1准确率，比RegNetY-16G高出0.3%，比DeiT-B高出0.1%。这些充满希望的结果强调了VMamba作为一个强大的基础模型的潜力，超越传统CNN模型和vision transformers的优势。

4.2 COCO数据集上的对象检测

设置

在本节中，我们使用MSCOCO 2017数据集[26]评估VMamba在对象检测方面的性能。我们的训练框架是建立在mmdetection库[2]上的，我们使用Mask-RCNN检测器遵循Swin[27]中的超参数。具体来说，我们使用了AdamW优化器，并对12和36个epoch的预训练分类模型(在ImageNet-1K上)进行了微调。对于VMamba-T/S/B，下降路径率分别设置为0.2%/0.2%/0.2% 。学习速率初始化为，并在第9和第11个epoch减少10倍。我们实现了多尺度的训练和批量为16的随机翻转。这些选择与已建立的对象检测评估实践相一致。

下降路径率（DropPath）是对 Dropout 的一种扩展，它不是仅将某个神经元的输出置零，而是随机删除整个路径（通常是某一层的所有连接）。在每个训练迭代中，每个路径的删除概率都为 0.2%。这样的设置使得网络的不同路径在训练过程中都有一定的几率被删除，从而增加了模型的鲁棒性，有助于防止过拟合。

结果

COCO的结果如表3所示。VMamba在COCO的box/mask平均精度(AP)上保持优势，无论采用的训练计划(12或36个epoch)。具体来说，VMamba-T/S/B模型通过12 epoch的微调，实现了46.5%/48.2%/48.5%的对象检测mAP，比Swin-T/S/B提高了3.8%/3.6%/1.6%的mAP，比ConvNeXt-T/S/B提高了2.3%/2.8%/1.5%的mAP。使用相同的配置，VMamba-T/S/B实现了42.1%/43.0%/43.1%的实例分割mIoU，分别优于Swin-T/S/B的结果2.8%/2.1%/0.8% mIoU和ConvNeXt-T/S/B的结果 2.0%/1.2%/0.7% mIoU。

此外，VMamba在36 epoch微调下，多尺度训练的优势依然存在，如表3所示。与Swin[28]、ConvNeXt[29]、PVTv2[49]、ViT[10](带适配器)等对等项相比，VMamba-T/S的性能更优越，目标检测的mAP值为48.5%/49.7%，实例分割的mIoU值为43.2%/44.0%。这些结果强调了VMamba在下游密集预测任务中的潜力。

4.3基于ADE20K的语义分割

设置

在Swin[28]之后，我们在预先训练的模型上构建一个UperHead[50]。利用AdamW优化器[30]，我们设定学习速率为。微调过程共跨越160k个迭代，批处理大小为16。默认的输入分辨率是512×512，我们还提供了使用640 × 640输入和多尺度(MS)测试的实验结果。

结果

结果如表4所示。VMamba模型再一次展示了卓越的精度，尤其是在分辨率为512 × 512的情况下，VMamba- t模型达到了47.3%的mIoU，在使用多尺度(MS)输入的情况下达到了48.3%的mIoU。这些分数超过了所有竞争对手，包括ResNet[19]、DeiT[45]、Swin[28]和ConvNeXt[29]。值得注意的是，这种优势可以扩展到VMambaS/B模型，即使在使用640 × 640输入时也是如此。

表4:使用UperNet[50]在ADE20K上的语义分词结果。我们用UperNet[50]在ADE20K数据集上评估了语义分割的性能。FLOPs 的计算是基于输入尺寸为 512 × 2048 或 640 × 2560，具体取决于裁剪尺寸。“SS”和“MS”分别表示单标度和多标度测试。

4.4分析实验

有效感受野

为了评估不同模型之间的有效接感受野(ERT)[32]，我们在图5中给出了一个比较分析。ERT度量模型输入与其输出之间的相关性。将输入大小为1024 × 1024的中心像素的ERF可视化，我们将VMamba与四种主要的视觉基础模型ResNet50[19]、ConvNeXt-T[29]、Swin-T[28]和DeiT-S [45] (ViT)在训练前和训练后阶段进行比较。图5的关键观察结果包括:1)只有DeiT (ViT)和VMamba显示全局ERT，而其他模型显示局部ERT，尽管它们理论上是全局的。值得注意的是，DeiT (ViT)模型产生了二次复杂度代价(见图6)。2)。与DeiT (ViT)使用注意机制均匀激活所有像素相比，VMamba激活所有像素，并特别强调十字形激活。跨扫描（Cross-Scan）模块的扫描机制确保中心像素受交叉像素的影响最大，将每个像素的长期依赖上下文优先于该像素的本地信息。3)有趣的是，VMamba最初在训练前只展示了一个局部的ERF。训练后将ERF转换为全局的，表示模型全局能力的一个自适应过程。我们相信，这种适应性过程有助于增强模型对图像的感知。这与DeiT形成了鲜明对比，VMamba在训练前和训练后都保持了几乎相同的ERT。

图5:ResNet50[19]、ConvNeXt-T[29]、Swin-T[28]、DeiT-S [45] (ViT))和VMamba-T的有效感受野(ERF)可视化。较大的ERF由分布更广泛的暗区表示。只有DeiT[45]和所提出的VMamba显示全局ERF。此可视化的灵感来自[32]。

图6:使用大小为224 × 224的图像作为输入训练主流模型进行评估，我们评估了各种主流模型的性能(a)和FLOPs (b)，这些模型使用大小为224 × 224的图像数据作为输入，不同的输入尺寸，范围从64 × 64到1024 × 1024。

输入尺寸调整（缩放）

我们继续进行输入缩放实验，测量ImageNet-1K和FLOPs上的top-1精度，如图6所示。在图6 (a)中，我们评估了主流模型(使用224×224输入大小训练)在各种图像分辨率(从64 × 64到1024 × 1024)下的推理性能。与同类产品相比，VMamba展示了在不同输入图像大小的情况下最稳定的性能。值得注意的是，随着输入大小从224 × 224增加到384 × 384，只有VMamba的性能有上升趋势(VMamba- s达到84%)，这突出了其对输入图像大小变化的鲁棒性。在图6 (b)中，我们评估了使用不同图像分辨率(也从64 × 64到1024 × 1024)的FLOPs。正如预期的那样，VMamba系列展现了与CNN模型一致的复杂的线性增长。VMamba的复杂性与Swin[28]等精心设计的vision transformers一致。然而，必须注意的是，只有VMamba才能实现全局有效感受野(ERF)。DeiT也具有全局ERF能力，其复杂性经历了二次增长。

5 结论

卷积神经网络(CNNs)和Vision Transformers(ViTs)是视觉表示学习的主要基础模型。CNNs在图像分辨率方面呈现线性的复杂度，而ViTs在拟合能力方面优于CNNs的二次复杂度。我们的调查显示，ViTs通过全局感受野和动态权值实现了卓越的视觉建模。基于此，我们提出了视觉状态空间模型(VMamba)，从状态空间模型中获得灵感，在不牺牲全局接收域的情况下实现线性的复杂度。为了解决方向敏感性问题，我们引入了交叉扫描模块(CSM)用于空间遍历，将非因果视觉图像转换为有序的patch序列。大量的实验证明VMamba在视觉任务中的表现很有前景，随着图像分辨率的提高，它具有明显的优势，超过了既定的基准。