《VMamba》论文笔记

原文链接：

[2401.10166] VMamba: Visual State Space Model (arxiv.org)

原文笔记：

What：

VMamba: Visual State Space Model

Why：

多年以来CNN和VIT作为视觉特征提取的主流框架

CNN具有模型简单，共享权重，计算效率高（线性复杂度）的优点，缺点是不具备全局感受野，能力不足VIT，不擅长处理多模态数据，并且模型已经比较老了，多种模型变体已经使卷积网络接近能力的上限

VIT基于Transformer，具有架构简单，全局感受野，动态权重(当提到"动态权重"时，通常是指注意力机制中用于分配不同输入元素之间关联度的权重。这些权重是根据输入数据的不同情况动态计算得出的。)的优点，缺点是计算效率不如卷积神经网络。（二次复杂度）

作者从最近很火的基于SSM模型的Mamba中获取灵感，集合了CNNs与VIT的优点提出了视觉状态空间模型VMamba该模型在不牺牲全局感受野的情况下实现了线性复杂度

Challenge：

1、作者想让这个视觉backbone在线性复杂度的情况下实现全局感受野，这是本研究的主要目的

2、如何将1D的SSM模型融合到2D的Vision任务当中作者将这一难题概括为

方向敏感问题：（然而，由于视觉数据的非因果性质，直接将这种策略应用于打成patch和拉伸后的图像不可避免地会导致感受野受限，因为无法估计与未扫描patch块的关系。我们将此问题称为“方向敏感”问题）

Idea：

为解决方向敏感问题作者提出交叉扫描模块(CSM)

（详见原文翻译部分）

Model：

(ps：时间有点紧，这篇整理的略显匆忙，但是这些基于Mamba的模型大差不差，连行文逻辑都很相似，将Mamba用作Vision领域的主要有两个工作一个是Vim一个就是本文VMamba，他们主要的区别都是在1D转2D这个操作上的区别，我将专门基于它们的区别写一篇)

原文翻译：

Abstract

卷积神经网络(CNNs)和Visual Transformer(ViTs)是视觉表示学习的两个最流行的基础模型。虽然 CNN 以线性复杂度在图像分辨率方面表现出显着的可扩展性，但 ViT 在拟合能力方面超过了它们，尽管它具有二次复杂度。仔细检查表明，ViTs通过结合全局感受野和动态权重实现了卓越的视觉建模性能。这一观察促使我们提出了一种新颖的架构，该架构继承了这些组件，同时提高了计算效率。为此，我们从最近引入的状态空间模型中获得灵感，提出了视觉状态空间模型(VMamba)，该模型在不牺牲全局接受域的情况下实现了线性复杂度。为了解决遇到的方向敏感问题，我们引入了交叉扫描模块(CSM)遍历空间域，并将任何非因果视觉图像转换为顺序补丁序列。大量的实验结果表明，随着图像分辨率的增加，VMamba不仅在各种视觉任务中表现出有希望的能力，而且在已建立的基准上也表现出更明显的优势。源代码可在 https://github.com/MzeroMiko/VMamba 获得。

1 Introduction

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

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

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

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

在不同的视觉任务上进行了大量的实验来验证VMamba的有效性。如图1所示，与Resnet[19]、ViT[10]和Swin[28]等基准视觉模型相比，VMamba模型在ImageNet-1K上表现出优越或至少具有竞争力的性能。我们还报告了下游密集预测任务的结果。例如，vambatiny /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模型，尽管通常仍然表现出优越的性能。然而，令人感兴趣的是VMamba本质上是基于Transformer体系结构的基础模型，它能够在FLOPs稳步增加的情况下获得与ViT相当的性能。

我们总结了本文的主要贡献：

• 我们提出了一个具有全局接受域和动态权值的视觉状态空间模型vamba，用于视觉表征学习。VMamba为视觉基础模型提供了一种新的选择，超越了cnn和vit的现有选择。
• 引入了交叉扫描模块(CSM)，以弥合一维阵列扫描和二维平面遍历之间的差距，促进S6扩展到视觉数据而不影响接收域。
• 我们展示了vamba在各种视觉任务(包括图像分类、对象检测和语义分割)上取得了令人满意的结果。这些发现强调了vamba作为一个强大的视觉基础模型的潜力。

2 Related Work

深度神经网络极大地促进了机器视觉感知的研究。视觉基础模型主要有两种，cnn[23,38,41,19,42]和vit[10,28,48,9,6,56]。最近，状态空间模型 (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]方面也取得了进展。

Vision Transformers(ViTs)是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 Preliminaries

状态空间模型。状态空间模型(SSMs)通常被认为是将刺激x(t)∈R^L映射到响应y(t)∈R^L的线性时不变系统。在数学上，这些模型通常被表述为线性常微分方程(ode) (Eq. 1)，其中参数包括A∈C^N ×N,  B, C∈C^N，状态大小为N，跳跃连接D∈C^1。

离散化。状态空间模型(ssm)作为连续时间模型，在与深度学习算法集成时面临着很大的挑战。为了克服这一障碍，离散化过程变得势在必行。离散化的主要目标是将ODE转化为离散函数。这种转换对于使模型与输入数据中隐含的底层信号的采样率保持一致至关重要，从而实现高效的计算操作[16]。考虑输入Xk∈R^L×D，在长度伟L的信号流内的采样向量如下[17]，常微分方程（Eq. 1）可以使用零阶保持规则离散如下：

选择性扫描机制。与主要关注线性时不变(LTI) SSM的流行方法不同，所提出的VMamba通过将选择性扫描机制(S6)[12]作为核心SSM算子而使自己与众不同。在S6中，矩阵B∈RB×L×N,C∈RB×L×N，∆∈RB×L×D由输入数据x∈RB×L×D导出。这意味着S6知道嵌入在输入中的上下文信息，从而确保该机制中的权重的动态性。

3.2 2D Selective Scan

尽管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 Overall Architecture

vamba - tiny的架构概述如图4 (a)所示。vamba通过使用类似于ViTs的干模块将输入图像划分为patch开始该过程，但没有进一步将补丁平坦化为1-D序列。这种修改保留了图像的二维结构，得到尺寸为H/4 × W/4 × C1的特征图。

然后vamba将几个VSS块堆叠在特征图上，保持相同的维度，构成“阶段1”。vamba中的分层表示是通过patch merge操作对“Stage 1”中的feature map进行降采样来构建的[27]。随后，更多的VSS块参与，导致输出分辨率为H/8 × W/8，形成“第二阶段”。重复此过程以创建分辨率分别为H/16 × W/16和H/32 × W/32的“阶段3”和“阶段4”。所有这些阶段共同构建了类似于流行的CNN模型[19,22,41,29,42]和一些vit[27, 48, 6, 56]的分层表示。所得到的体系结构可以作为具有类似需求的实际应用中其他视觉模型的通用替代品。

我们开发了三种不同的vamba，即vamba - tiny, vamba - small和vamba - ase(分别称为vamba - t, vamba - s和vamba - b)。表1概述了详细的体系结构规范。所有模型的FLOPs都使用224 × 224的输入尺寸进行评估。在未来的更新中将引入其他架构，例如大型模型。

3.3.2 VSS Block

VSS块的结构如图4 (b)所示。输入经过初始线性嵌入层，输出分成两个信息流。在进入核心SS2D模块之前，一个流通过一个3 × 3深度卷积层，然后是一个Silu激活函数[37]。SS2D的输出经过一个层规范化层，然后加入到另一个信息流的输出中，这个信息流经过了一个Silu激活。这种组合产生了VSS块的最终输出。与VisionTransformer不同，由于其因果性，我们避免在VMamba中使用位置嵌入偏差。我们的设计不同于典型的视觉转换器结构，它在一个块中采用了如下的操作顺序:Norm → attention → Norm → MLP，然后丢弃MLP操作，因此VSS块比VIT块更浅，这是我们能够以相似的总模型深度预算堆叠更多块。

4 Experiment

略