VMamba:视觉状态空间模型

摘要

受最近提出的状态空间模型启发，我们提出了视觉状态空间模型，在不牺牲感受野的情况下实现线性复杂度。为了解决遇到的方向感知问题，我们提出了交叉扫描模块（CSM）以遍历空间域并将任何非因果视觉图像转换为有序补丁序列。广泛实验结果说明VMamba不仅在多种视觉感知任务上取得优秀的成果，在不断提升图像分辨率训练表现上也展现了明显提升。

VMamba在有效降低注意复杂度的关键概念来自选择性扫描状态空间模型（S6），原本被提出用来解决NLP问题。和传统注意力计算方法不同，S6使1-D矩阵内的每个元素来和先前扫描的样本交互，通过一个压缩的隐状态，有效将平方复杂度降低到线性。

然而，由于视觉信息是非因果的，直接在一个切分为patch并展平的图像上应用这样的策略会导致受限的感受野，因为和没被扫描到的patch间的关系难以被评估。我们称这个问题为“方向感知”问题并且提出通过新提出的交叉扫描模块来解决这一问题。不同于翻转特征图的空域以无方向模式（行方向或列方向），CSM采用一种四路扫描策略即从整个特征图的四个角到相反的位置。

总结我们的贡献如下：

• 我们提出了VMamba,有全局感受野的视觉状态空间模型
• 提出CSM弥补1-D和2-D平面之间的差距
• 没有花里胡哨的，我们展示VMamba在诸如图片分类，目标检测和语义分割任务上取得了很大的成功。

相关工作

状态空间模型

这些模型主要关注于状态空间模型如何应用在长程和因果数据像语言语音，语言理解，文本推理，像素级1-D图像分类上。S4ND是第一个将状态空间机制应用在视觉任务上的工作，但是不能以依赖于输入的方式有效地捕获图像信息。

方法

准备

状态空间模型

见【论文翻译】Mamba 中的状态空间模型背景-CSDN博客

离散化

见【论文翻译】MODELING SEQUENCES WITH STRUCTURED STATE SPACES 附录A.1-CSDN博客

VMamba在这里使用ZOH零阶保持离散化，但是$\overline{B}$使用一阶泰勒近似为
$$\overline{B}=(e^{\Delta A}-I)A^{-1}B\approx (\Delta A)(\Delta A{)}^{-1}\Delta B$$

选择扫描机制

在S6，参数$B,C,\Delta$是数据依赖的，确保了权重的动态性。

2D 选择扫描

S6因果的处理输入数据，因此只能捕捉到已经扫描到的部分的信息。天然适应NLP任务包括时许信息，但对于非因果的数据例如图片，图，集等等设置了挑战。一个直接的解决这个问题的办法是从两个方向扫描数据，即前向和后向，允许它们在不增加计算复杂度的情况下互相补充感受野。

除了有非因果的天然性质，图像和文本间的差别还包括图像包含2D的空间信息（例如局部信息和全局信息）。为了处理这个问题，S4ND提出用卷积和直接扩展核从1-D到2-D通过外积来重新公式化SSM。然而，这样的调整使得权重不能变成动态的（即输入依赖的）导致基于内容的数据建模的能力损失。因此，我们选择保留动态权重通过坚持选择性扫描方法，不幸的是这阻碍我们遵循卷积操作。

为了让解决这个问题，我们提出了图二所示的交叉扫描魔魁啊。我们选择展开图像块沿着行和列到需略（扫描扩展），然后从四个不同方向处理扫描：上-左到下-右，下-右到上-左，上-右到下-左，下-左到上-右。以这种方式，任何像素如图二的中心像素，整合了来自不同方向的其他像素。我们之后将每个序列reshape称单张图片，然后所有的序列被拼接成一个新的如图三所描述。（扫描拼接）

和CSM结合的S6，被称作S6块，是视觉状态空间（VSS）块的核心元素，构成了VMamba的基础块（在下一孝介进一步介绍）。我们强调在保持全局感受野情况下S6继承了选择扫描机制的线性复杂度。

VMamba 模型

整体结构

如图4（a）所示，是VMamba-Tiny的概览。VMamba首先使用一个stem模块把输入图像分割为块，和ViTs相似不需要进一步展平patch到1-D序列。这个调整保持了图片的2D结构，使特征图有$\frac{H}{4}\times \frac{W}{4}\times C_1$

VMamba之后在特征图上堆叠一些VSS块，保持相同的维度，构成阶段1。VMamba中的层次表示是通过pach合并操作对“阶段1”中的特征图进行下采样来构建的。其后，更多的VSS块包含进来构成阶段2，输出分辨率为$\frac{H}{8}\times \frac{W}{8}$.这个过程重复来构成阶段3和阶段4，分辨率分别为$\frac{H}{16}\times \frac{W}{16}$和$\frac{H}{32}\times \frac{W}{32}$。所有这些阶段共同构成类似于流行CNN和一些ViT的层次表示。

我们在三个不同规模下实现了VMamba，VMamba-Tiny, VMamba-Small, and VMamba-Base。所有模型的FLOPs通过使用$224\times 224$​输入评估。

VSS块

图四b展示了VSS块的结构，输入经过一个初始线性嵌入层，其输出分为两股信息流。一个经过一个$3\times 3$的深度卷积层，在进入核心SS2D模块前经过一个Silu激活函数。SS2D的输出经过一个层归一化，之后把输出加到另一个信息流中。这组成了VSS块的最终输出。

不同于vision transformers我们避免在VMamba使用位置嵌入偏置由于其因果性。我们的设计偏离了典型的视觉transformer结构，其在块中应用了下面的操作顺序：归一化->注意力->归一化->MLP。我们抛弃掉MLP操作。因此，VSS块比ViT块浅一点，这允许我们在同样模型深度情况下堆叠更多的块。

实验

在这个部分，我们完成了一系列实验评估和比较VMmaba和流行模型，包括CNN和视觉transformers。我们的评估基于多种任务，包括ImageNet-1K图片分类，COCO目标检测和ADE20K语义分割。随后，我们深入分析实验，以深入了解VMamba的架构。

分析实验

有效感受野

为了评估不同模型的有效感受野，我们在图5展示了一个比较分析。ERT测量了模型输入对模型输出的影响。在输入尺寸1024X1024条件下可视化中心像素ERF。我们将其与四个主流的视觉基础模型作比较：从图五，我们得出关键结论1）只有DeiT和VMamba有全局ERF 2）不同于DeiT平等均匀地激活所有像素，VMamba激活所有像素但是显著强调了交叉型的激活。交叉扫描模块扫描机制使得中心像素受交叉方向像素影响最大，对于每个像素，其长程内容优先于局部信息。3）有趣的是，VMamba在训练前初始化仅仅一个局部ERF而在训练后变为全局，注意到模型全局能力的自适应过程。我们认为这个自适应过程增强了模型对图片的感知。和DeiT在训练前后都保持几乎一致地ERF不同。

输入尺度

我们继续进行输入尺度变化实验，如图6所示，VMamba实现了跨输入图片尺寸地最稳定表现，值得注意的是，在从224x224到384x384，只有VMamba表现出一个上升趋势，展示了其对输入尺寸变化地鲁棒性，VMamba实现了比肩CNN地线性复杂度增长，而同样有全局感受野地DeiT表现出了平方复杂度增长。

总结

为了解决方向感知问题，我们提出里交叉扫描模块来空间翻转，将非因果的视觉图像转换为顺序patch序列。大量的实验表明，VMamba在多种视觉任务上取得可观表现，当图片分辨率上升时，有显著的优势。