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$。
h ′