VMamba: Visual State Space Model

VMamba: Visual State Space Model

• 论文地址: https://arxiv.org/abs/2401.10166
• 项目地址: https://github.com/MzeroMiko/VMamba

Abstract

卷积神经网络(CNN)和视觉Transformer(ViT)是视觉表征学习中最流行的两种基础模型。CNN在线性复杂度和图像分辨率方面表现出显著的可扩展性，而ViT在拟合能力方面超越了CNN，尽管它要与二次复杂度竞争。通过对全局感受野和动态权重的结合，ViT实现了卓越的视觉建模性能。这一观察结果促使我们提出一种新的架构，该架构继承了这些组件，同时提高了计算效率。为此，我们从最近引入的状态空间模型中汲取灵感，提出了在不牺牲全局感受野的情况下实现线性复杂性的视觉状态空间模型(VMamba)。为了解决遇到的方向敏感问题，我们引入了交叉扫描模块(CSM)来遍历空间域并将任何非因果视觉图像转换为有序patch序列。大量的实验结果证实，VMamba不仅在各种视觉感知任务中表现出有希望的能力，而且随着图像分辨率的增加，它比现有基准显示出更明显的优势。

1 Introduction

深度基础模型的两个主要类别，即卷积神经网络(CNN)[38, 19, 22, 29, 42]和视觉Transformer(ViT)[10, 28, 45, 56]，已经广泛应用于各种视觉任务中。虽然两者在计算表达性视觉表征方面都取得了显著的成功，但ViT总体上表现出优于CNN的性能，这可能归因于全局感受野和注意力机制促进的动态权重。

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

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

然而，由于视觉数据的非因果性，直接将这种策略应用于patchified和flattened图像将不可避免地导致受限制的感受野，因为无法估计与未扫描patch的关系。我们将这个问题称为“方向敏感”问题，并建议通过新引入的交叉扫描模块(CSM)来解决这个问题。CSM不是以单向模式(列向或行向)遍历图像特征映射的空间域，而是采用四向扫描策略，即从特征映射的四个角到相反的位置(见图2 (b))。该策略确保了特征映射中的每个元素都集成了来自不同方向的所有其他位置的信息，从而在不增加线性计算复杂度的情况下呈现全局感受野。

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

我们总结了以下贡献：

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

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

• 没有花里胡哨，我们表明VMamba在各种视觉任务中取得了令人满意的结果，包括图像分类，目标检测和语义分割。这些发现强调了VMamba作为一个强大的视觉基础模型的潜力。

2 Related Work

深度神经网络极大地推动了机器视觉感知的研究。主要有两种流行的视觉基础模型，即CNN[23、38、41、19、42]和ViT[10、28、48、9、6、56]。近年来，状态空间模型(State Space Models, SSM)[12,34,47]取得了成功，说明了它们在高效长序列建模中的有效性，这引起了NLP和CV社区的广泛关注。我们的研究坚持这条路线，并提出了VMamba，一种基于SSM的视觉领域数据建模架构。VMamba与CNN和ViT一起为社区提供了另一种基础模型。

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

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

状态空间模型(SSM)是最近提出的模型，作为状态空间转换引入深度学习[16,15,39]。受控制系统中的连续状态空间模型的启发，结合HiPPO[13]初始化，LSSL[16]展示了处理远程依赖问题的潜力。然而，由于状态表示对计算量和内存的要求过高，LSSL在实际应用中并不可行。为了解决这个问题，S4[15]提出将参数归一化为对角结构。此后，不同结构的结构状态空间模型层出不穷，如复杂对角结构[17,14]、多输入多输出支持[39]、对角分解加低秩运算[18]、选择机制[12]等。然后将这些模型集成到大型表示模型中[34,33,11]。

这些模型主要关注如何将状态空间模型应用于语言和语音等远程和随机数据，如语言理解[33,34]，基于内容的推理[12]，像素级的一维图像分类[15]，在视觉识别方面很少受到重视。与我们最相似的工作是S4ND[35]。S4ND是第一个将状态空间机制应用到视觉任务中的工作，并显示了其性能可能与ViT竞争的潜力[10]。然而，S4ND对S4模型进行了简单的扩展，无法以依赖输入的方式有效地捕获图像信息。我们证明了通过mamba[12]引入的选择性扫描机制，所提出的VMamba能够匹配现有流行的视觉基础模型，如ResNet[19]、ViT[10]、swin[27]和convnext[29]，显示了VMamba作为强大基础模型的潜力。

3 Method

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

3.1 准备工作

状态空间模型。状态空间模型(SSM)通常被认为是线性时不变系统，将刺激$x(t)\in {\mathbb{R}}^L$映射到响应$y(t)\in {\mathbb{R}}^L$。在数学上，这些模型通常被表述为线性常微分方程(ODE) (公式1)，其中参数包括$A\in {\mathbb{C}}^{N\times N},B,C\in {\mathbb{C}}^N$，对于状态大小$N$，跳跃连接$D\in {\mathbb{C}}^1$。
$$\begin{array}{rl}& h^{\prime }(t)=Ah(t)+Bx(t)\\ & y(t)=Ch(t)+Dx(t)\end{array}\text{\hspace{1em}(1)}$$
离散化。状态空间模型(SSM)作为连续时间模型，在与深度学习算法集成时面临着很大的挑战。为了克服这一障碍，离散化过程变得势在必行。

离散化的主要目标是将ODE转化为离散函数。这种转换对于使模型与输入数据中隐含的底层信号的采样率保持一致至关重要，从而实现高效的计算操作[16]。考虑输入$x_k\in {\mathbb{R}}^{L\times D}$，在长度为L的信号流内的采样向量如下[17]，ODE (公式1)可以使用零阶保持规则离散如下：
$$\begin{array}{rl}{h}_k& =\bar{A}{h}_{k-1}+\bar{B}{x}_k,\\ y_k& =\bar{C}{h}_k+\bar{D}{x}_k,\\ \bar{A}& =e^{\Delta A},\\ \bar{B}& =\left(e^{\Delta A}-I\right)A^{-1}B,\\ \bar{C}& =C\end{array}\text{\hspace{1em}(2)}$$
其中，$B,C\in {\mathbb{R}}^{D\times N}$和$\Delta \in {\mathbb{R}}^D$。在实践中，下面[12]，我们使用一阶泰勒级数来细化$\bar{B}$的近似：
$$\bar{B}=\left(e^{\Delta A}-I\right)A^{-1}B\approx (\Delta A)(\Delta A{)}^{-1}\Delta B=\Delta B\text{\hspace{1em}(3)}$$
选择性扫描机制(Selective Scan Mechanism)。与主要关注线性时不变(LTI) SSM的流行方法不同，所提出的VMamba通过将选择性扫描机制(S6)[12]作为核心SSM算子而使自己与众不同。在S6中，矩阵$B\in {\mathbb{R}}^{B\times L\times N}$, $C\in {\mathbb{R}}^{B\times L\times N}$，$\Delta \in {\mathbb{R}}^{B\times L\times D}$由输入数据$x\in {\mathbb{R}}^{B\times L\times D}$导出。这意味着S6知道嵌入在输入中的上下文信息，从而确保该机制中的权重的动态性。

3.2 二维选择性扫描

尽管S6具有鲜明的特点，但它对输入数据进行因果处理，因此只能捕获数据扫描部分内的信息。这自然使S6与涉及时间数据的NLP任务保持一致，但在适应非因果数据(如图像、图形、集合等)时提出了重大挑战。这个问题的一个直接解决方案是沿着两个不同的方向(即向前和向后)扫描数据，允许它们在不增加计算复杂性的情况下补偿彼此的感受野。

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

为了解决这个问题，我们提出了如图2所示的交叉扫描模块(CSM)。

我们选择将图像patch沿着行和列展开成序列(扫描展开)，然后沿着四个不同的方向进行扫描：左上到右下，右下到左上，右上到左下，左下到右上。这样，任何像素(如图2中的中心像素)都可以集成来自不同方向的所有其他像素的信息。然后，我们将每个序列重塑为单个图像，并将所有序列合并为一个新的序列，如图3所示(扫描合并)。

S6与CSM的集成称为S6块，它是构建可视状态空间(VSS)块的核心元素，VSS块构成了VMamba的基本构建块(下一小节将进一步详细介绍)。我们强调，S6块继承了选择性扫描机制的线性复杂性，同时保留了全局感受野，这与我们构建这种视觉模型的动机是一致的。

3.3 VMamba Model

3.3.1 总体架构

VMamba-Tiny的架构概述如图4 (a)所示。

VMamba通过使用类似于ViT的干模块将输入图像划分为patch开始该过程，但没有进一步将patch平坦化为1-D序列。这种修改保留了图像的二维结构，得到尺寸为H/4 × W/4 × C₁的特征图。然后VMamba将几个VSS块堆叠在特征图上，保持相同的维度，构成“阶段1”。VMamba中的分层表示是通过patch merge操作对“阶段1”中的feature map进行降采样来构建的[27]。随后，涉及更多的VSS块，导致输出分辨率为H/8 × W/8，形成“阶段2”。重复此过程以创建“阶段3”和“阶段4”，分辨率分别为H/16 × W/16和H/32 × W/32。所有这些阶段共同构建了类似于流行的CNN模型[19,22,41,29,42]和一些ViT[27, 48, 6, 56]的分层表示。所得到的体系结构可以作为具有类似需求的实际应用中其他视觉模型的通用替代品。我们开发了三种不同的VMamba，即VMamba -Tiny， VMamba-Small和VMamba-Base(分别称为VMamba-T，VMamba-S和VMamba-B)。

表1概述了详细的体系结构规范。所有模型的FLOPs都使用224 × 224的输入尺寸进行评估。在未来的更新中将引入其他架构，例如大型模型。

3.3.2 VSS模块

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

与视觉Transformer不同，由于其因果性，我们避免在VMamba中使用位置嵌入偏差。我们的设计与典型的视觉Transformer结构不同，后者采用以下操作顺序：Norm → attention → Norm → MLP 在每个模块中，并且丢弃MLP操作。因此，VSS块比ViT块更浅，这使我们能够以相似的总模型深度预算堆叠更多块。

4 Experiment

在本节中，我们执行一系列实验来评估和比较VMamba与流行的模型，包括CNN和视觉Transformer。我们的评估跨越了不同的任务，包括ImageNet-1K上的图像分类、COCO上的目标检测和ADE20K上的语义分割。随后，我们深入分析实验，以更深入地了解VMamba的架构。

4.1 ImageNet-1K上的图像分类

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

结果 表2总结了ImageNet-1K上的结果，将VMamba与流行的CNN模型和视觉Transformer进行了比较。对比发现，在FLOPs相似的情况下，VMamba-T的性能达到82.2%，比RegNetY-4G高2.2%，比DeiT-S高2.4%，比Swin-T高0.9%。值得注意的是，VMamba的性能优势在小型模型和基本模型中仍然存在。例如，在小规模测试中，VMamba-S的准确率达到了最高的83.5%，比RegNetY-8G高1.8%，比Swin-S高0.5%。同时，VMamba-B实现了83.2%的前1准确率，超过RegNetY-16G 0.3%和DeiT-B 0.1%。这些有希望的结果强调了VMamba作为强大基础模型的潜力，扩展了其超越传统CNN模型和视觉Transformer的优势。

4.2 COCO上的目标检测

设置 在本节中，我们使用MSCOCO 2017数据集[26]评估了所提出的VMamba在目标检测方面的性能。我们的训练框架是建立在mmdetection库[2]上的，我们坚持使用Mask-RCNN检测器的Swin[27]中的超参数。具体来说，我们使用了AdamW优化器，并对预训练的分类模型(在ImageNet-1K上)进行了12和36个epoch的微调。VMamba-T/S/B的掉落路径率分别设置为0.2%/0.2%/0.2%。学习率初始化为1 × 10⁻⁴，并在第9和第11次迭代时降低10倍。我们实现了批量大小为16的多尺度训练和随机翻转。这些选择与目标检测评估的既定实践保持一致。

结果 COCO的结果总结于表3。无论采用何种训练计划(12或36次)，VMamba在COCO的盒/掩模平均精度(AP)方面都保持优势。其中，VMamba-T/S/B模型在12epoch的微调调度下，实现了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实现的实例分割mIoU分别为42.1%/43.0%/43.1%，分别优于Swin-T/S/B 2.8%/2.1%/0.8% mIoU，优于ConvNeXt-T/S/B 2.0%/1.2%/0.7% mIoU。

此外，在多尺度训练的36epoch微调时间表下，VMamba的优势仍然存在，如表3所示。与Swin[28]、ConvNeXt[29]、PVTv2[49]和ViT[10]（带Adapter）等对应产品相比，VMamba-T/S表现出卓越的性能，在对象检测上实现了48.5%/49.7%的mAP，在实例分割上实现了43.2%/44.0%的mIoU。这些结果强调了VMamba在下游密集预测任务中的潜力。

4.3 ADE20K上的语义分割

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

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

4.4 分析实验

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

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

5 Conclusion

卷积神经网络(CNN)和视觉Transformer(ViT)是视觉表征学习的主要基础模型。CNN在图像分辨率方面表现出线性复杂性，而ViTs在拟合能力方面表现出色，尽管具有二次复杂度。我们的研究表明，ViTs通过全局感受野和动态权重实现了卓越的视觉建模。受此启发，我们提出了视觉状态空间模型(VMamba)，从状态空间模型中汲取灵感，在不牺牲全局感受野的情况下实现线性复杂性。为了解决方向敏感性问题，我们引入了用于空间遍历的交叉扫描模块(CSM)，将非因果视觉图像转换为有序的补丁序列。大量的实验证明VMamba在视觉任务中的表现很有希望，随着图像分辨率的提高，它的优势明显，超过了既定的基准。