VMamba论文精读笔记

论文链接：[2401.10166] VMamba: Visual State Space Model

Institution: 中国科学院大学（UCAS）, 华为, 鹏城实验室

代码链接：GitHub - MzeroMiko/VMamba: VMamba: Visual State Space Models，code is based on mamba

摘要

在计算机视觉领域，设计计算高效的网络架构始终是一项持续的需求。在本文中，我们将语言模型 Mamba（一个状态空间模型）适配为 VMamba，一种具有线性时间复杂度的视觉骨干网络（vision backbone）。VMamba 的核心组件是由 视觉状态空间（Visual State-Space, VSS）块 和 二维选择性扫描（2D Selective Scan, SS2D）模块 组成的堆栈。通过沿着++四种扫描路径++进行遍历，SS2D 弥合了 1D 选择性扫描有序特性与 2D 视觉数据非序列化结构之间的差距，从而有助于从多个来源和视角收集上下文信息。

基于 VSS 块，我们开发了一个 VMamba 系列架构，并通过一系列架构上的和实现上的改进来加速其性能。广泛的实验表明，VMamba 在各种视觉感知任务中展现了令人期待的性能，尤其是与现有的基准模型相比，VMamba 在输入规模扩展效率上表现出了显著优势。

一、引言

视觉表示学习作为计算机视觉中的一个基础研究领域，在深度学习时代取得了显著进展。为了表示视觉数据中的复杂模式，两类主要的骨干网络被提出并广泛应用于各种视觉任务中：卷积神经网络（CNNs） [49, 27, 29, 53, 37] 和 视觉Transformer（ViTs） [13, 36, 57, 66]。与 CNNs 相比，ViTs 由于整合了自注意力机制 [58, 13]，通常在大规模数据上展现出更强的学习能力。然而，自注意力机制的计算复杂度相对于 token 数目的二次增长，在处理大空间分辨率的下游任务时会带来显著的计算开销。为了解决这一问题，人们做出了大量努力来提高注意力计算的效率 [54, 36, 12]。然而，现有方法要么限制了有效感受野的大小 [36]，要么在各种任务上表现出显著的性能下降 [30, 60]。

这一现状促使我们开发一种用于处理视觉数据的新型架构，同时保留原生自注意力机制的固有优势，例如全局感受野 和动态加权参数 [23]。最近，在自然语言处理（NLP）领域，Mamba [17] 作为一种创新的状态空间模型（State Space Model, SSM） [17, 43, 59, 71, 48]，已被证明是长序列建模的一个有效方法，且具有线性计算复杂度。受这一进展的启发，我们提出了 VMamba，一种集成基于 SSM 块的视觉骨干网络，旨在实现高效的视觉表示学习。

然而，Mamba 的核心算法，即++并行化选择扫描操作，本质上是为处理一维序列数据而设计的。当适配到视觉数据时，视觉数据缺乏固有的顺序排列++ ，这带来了挑战。为了解决这一问题，我们提出了 ++二维选择性扫描（2D Selective Scan, SS2D）++ ++，这是一种专为空间域遍历设计的四向扫描机制++ 。与自注意力机制（图1 (a)）相比，SS2D 确保每个图像 patch 仅通过沿其对应扫描路径计算的压缩隐藏状态获取上下文信息（图1 (b)），从而将计算复杂度从二次降低为线性。

（a）self-Attention :传统自注意力机制Self-Attention）直接为目标图像 patch（红色框，查询 patch）与其他所有图像 patch 建立全局关联（黄色箭头表示不同 patch 的交互）。自注意力机制通过计算**每个 patch 与所有其他 patch 的关系（即全局交互）**来提取上下文信息。这样可以有效捕获全局感受野中的信息。

（b）SS2D: 通过四种扫描路径（蓝色箭头：从左到右、从右到左；黄色箭头：从上到下、从下到上）来遍历图像。红色框中目标图像 patch（查询 patch）的上下文信息仅通过沿其对应扫描路径计算的压缩隐藏状态获取，而不是与所有其他 patch 做直接交互。 通过这种方式，SS2D 在提取上下文信息的同时显著降低了计算复杂度。

基于 VSS（视觉状态空间）块，我们开发了一系列 VMamba 架构（即 VMamba-Tiny/Small/Base），并通过架构改进和实现优化提升了其性能。与基于 CNNs（如 ConvNeXt [37]）、ViTs（如 Swin [36]、HiViT [66]）和 SSMs（如 S4ND [44]、Vim [69]）的基准视觉模型相比，VMamba 在 ImageNet-1K [9] 上的各种模型规模下均实现了更高的图像分类准确率。具体来说，VMamba-Base 在 ImageNet-1K 上取得了 83.9% 的 top-1 准确率，比 Swin 高出 0.4%，且吞吐量比 Swin 提高了显著的 40% 以上（646 vs. 458）。VMamba 的优势还扩展到多个下游任务，其中 VMamba-Tiny/Small/Base 在 COCO [33] 上的目标检测中（1× 训练日程）分别实现了 47.3%/48.7%/49.2% 的 mAP，相较于 Swin 提高了 4.6%/3.9%/2.3%，相较于 ConvNeXt 提高了 3.1%/3.3%/2.2%。至于在 ADE20K [68] 上的单尺度语义分割任务中，VMamba-Tiny/Small/Base 分别实现了 47.9%/50.6%/51.0% 的 mIoU，相较于 Swin 提升了 3.4%/3.0%/2.9%，相较于 ConvNeXt 提升了 1.9%/1.9%/1.9%。

此外，与 ViT 模型在输入 token 数量增加时计算复杂度呈二次增长不同，VMamba 在保持可比性能的同时，其 FLOPs 呈线性增长。这展示了其在输入可扩展性方面的最先进表现。

什么是FLOPs?

FLOPs 是 Floating Point Operations per Second 的缩写，指的是 **每秒浮点运算次数。**在深度学习中，FLOPs 通常用来表示模型完成一次前向传播、反向传播所需的浮点运算总量，是衡量模型计算效率的重要指标。

FLOPs 越低，说明模型越高效，可能适合部署到更多场景中。

本文贡献总结

提出 VMamba：一种基于 SSM 的视觉骨干网络，用于实现具有线性时间复杂度的视觉表示学习。我们采用了一系列架构和实现优化来提升 VMamba 的推理速度。
引入二维选择性扫描（SS2D）：通过连接一维数组扫描和二维平面遍历，使选择性 SSM 能够扩展到处理视觉数据。
多任务性能优越：VMamba 在图像分类、目标检测和语义分割等各种视觉任务中表现出色，同时展现了在输入序列长度方面的高适应性，计算复杂度呈线性增长。

二、相关工作

卷积神经网络（CNNs）

自 AlexNet [31] 提出以来，人们在各种视觉任务中对基于 CNN 的模型进行了大量改进，致力于提升其建模能力 [49, 52, 27, 29] 和计算效率 [28, 53, 64, 46]。例如，引入了深度可分离卷积（depth-wise convolution）[28] 和可变形卷积（deformable convolution）[5, 70] 等高级操作，来增加 CNN 的灵活性和效率。最近，受到 Transformer [58] 成功的启发，现代 CNN 模型 [37] 通过将长距离依赖关系 [11, 47, 34] 和动态权重 [23] 融入其设计中，展现了令人期待的性能。

视觉 Transformer（ViTs）

作为开创性工作，ViT [13] 探索了基于标准 Transformer 架构的视觉模型的有效性，并强调了大规模预训练对图像分类性能的重要性。为了减少 ViT 对大数据集的依赖，DeiT [57] 提出了教师-学生蒸馏策略，通过从 CNN 向 ViT 转移知识的方法，突出了++归纳偏置（inductive bias）在视觉感知中的重要性。++ 在此基础上，后续研究提出了分层 ViT [36, 12, 61, 39, 66, 55, 6, 10, 67, 1]。

什么是归纳偏置Inductive bias?

指的是模型在学习过程中所采用的假设或先验知识，这些假设帮助模型从有限的训练数据中推断出更广泛的模式或规律。

在视觉感知中，常见的归纳偏置类型包括：

1. 空间局部性（Spatial Locality）

定义：图像中的信息往往是局部相关的，即一个像素的值通常与其周围的像素密切相关。

应用：卷积神经网络（CNN）通过局部感受野和卷积核实现了这一归纳偏置。卷积核只关注局部区域（如 3×3或 5×5），从而有效地提取局部特征。

效果：这种归纳偏置使 CNN 更适合处理图像数据，因为它们能高效地捕获局部模式。

2. 平移不变性（Translation Invariance）

定义：图像中的对象无论如何平移，其本质特征保持不变。

应用：CNN 中的卷积操作和池化操作天然支持平移不变性。例如，卷积核在图像上滑动（共享权重），能够检测特定特征的位置无关性。

效果：这一偏置提高了模型对位置信息的鲁棒性，使其在不同位置检测相同对象。

3. 层次结构（Hierarchical Structure）

定义：图像中的信息通常具有分层结构，例如低级特征（边缘、角点）组合成中级特征（纹理、形状），最终构成高级特征（对象）。

应用：深度神经网络中的多层结构（尤其是 CNN）反映了这种归纳偏置。底层网络提取局部边缘和纹理，高层网络提取全局对象信息。

效果：这种分层归纳偏置使模型能够逐步学习从简单到复杂的视觉特征。

4. 对称性（Symmetry）

定义：许多视觉数据中的特征具有对称性，例如水平对称或旋转对称。

应用：通过数据增强（如翻转、旋转）或特定模型设计（如旋转等变卷积），模型能够有效捕获对称性特性。

效果：这种归纳偏置帮助模型减少对数据规模的需求，同时提升其对对称模式的理解能力。

5. 稀疏连接（Sparse Connectivity）

定义：在视觉中，一个特定的像素只与其局部邻域有直接关联，而不是与整个图像相关联。

应用：CNN 的局部连接结构通过稀疏连接（而非全连接）实现了这一偏置。

效果：稀疏连接大幅降低了计算复杂度，使得 CNN 在处理高维输入（如图像）时更加高效。

6. 归纳层次感（Inductive Hierarchy of Objects）

定义：图像中的物体或场景往往可以被分解成若干子结构。

应用：分层模型（如分层 ViT）通过以不同尺度提取特征，捕获从局部到全局的信息。

效果：这种偏置帮助模型理解复杂场景中的多级结构。

另一研究方向聚焦于提高自注意力机制（即 ViT 的核心）的计算效率。例如：

线性注意力（Linear Attention） [30] 将自注意力重新表述为核特征映射的线性点积，通过矩阵乘法的结合律（associativity property）将计算复杂度从二次降至线性。
GLA（Global Linear Attention） [65] 引入了一种硬件高效的线性注意力变体，平衡了内存移动和并行化能力。
RWKV [45] 利用线性注意力机制结合了可并行化的 Transformer 训练与高效的循环神经网络（RNN）推理能力。
RetNet [51] 添加了门控机制（gating mechanism），使计算路径支持并行化，提供了循环计算的另一种替代方案。
RMT [15] 将时序衰减机制（temporal decay mechanism）扩展到了空间域，用于视觉表示学习。

状态空间模型（SSMs）

尽管 ViT 架构在视觉任务中被广泛采用，但由于++自注意力的二次复杂度++ （尤其是在处理长序列输入时，如高分辨率图像），它面临着显著的扩展性挑战。在提升扩展效率 [8, 7, 45, 51, 41] 的努力中，状态空间模型（State Space Models, SSMs） 作为 Transformer 的有力替代方案，吸引了大量研究关注。

Gu 等人 [21] 展示了 SSM 模型在处理长距离依赖关系时的潜力，使用了 HiPPO 初始化方法 [18]。为了提升其实际可行性，S4 [20] 提出了将参数矩阵归一化为对角结构的改进方法。此后，出现了多种结构化 SSM 模型，每种模型都提供了独特的架构改进，例如：

复数对角结构（complex-diagonal structures） [22, 19]；
支持多输入多输出（multiple-input multiple-output, MIMO） [50]；
对角加低秩分解（diagonal plus low-rank decomposition） [24]；
选择机制（selection mechanisms） [17]。

这些进步还被集成到更大的表示学习模型中 [43, 41, 16]，进一步突出了结构化状态空间模型在各种应用中的多功能性和可扩展性。虽然这些模型主要针对文本和语音等长距离和序列数据，但应用 SSMs 于具有二维结构的视觉数据的研究仍然较少。

三、初步概念

状态空间模型（SSMs）的公式化

状态空间模型（SSMs）起源于卡尔曼滤波器 [32]，是一种线性时不变（LTI）系统，通过隐状态将输入信号映射到输出响应。具体来说，连续时间 SSM 可以用以下线性常微分方程（ODE）表示：

其中，、、和是权重参数。

SSM 的离散化

为了能够集成到深度模型中，连续时间的 SSM 必须先经过离散化处理。具体来说，对于时间区间，在时隐状态变量的解析解可以表示为：

通过采样时间尺度参数 Δ（即），隐状态可以被离散化为：

其中，是对应的离散步长区间。值得注意的是，这种公式化近似于零阶保持（ZOH）方法的结果，该方法在 SSM 基模型的文献中经常使用（详见附录 A 的详细证明）。

选择性扫描机制

为了克服线性时不变 SSMs（公式 1）在捕获上下文信息方面的局限性，Gu 等人 [17] 提出了一个新颖的 SSM 参数化方法，该方法结合了依赖于输入的选择机制（称为 S6）。然而，对于选择性 SSMs，时变权重参数给隐状态的有效计算带来了挑战，因为卷积不能处理动态权重，导致其无效。然而，由于可以推导出公式 3 中的的递归关系，因此响应仍然可以使用关联扫描算法 [2, 42, 50] 高效计算，该算法具有线性复杂度（有关详细解释，请参见附录 B）。

四、VMamba:视觉状态空间模型

4.1 网络架构

我们开发了 VMamba 模型，分为三个规模：Tiny、Small 和 Base（分别称为 VMamba-T、VMamba-S 和 VMamba-B）。VMamba-T 的架构概述如图 3 (a) 所示，详细配置见附录 E。

输入图像首先通过一个初始模块被划分为多个小块，生成一个空间维度为的二维特征图。为了创建具有,和分辨率的分层表示，采用了多个网络阶段，而不加入额外的位置信息嵌入。

具体来说，每个阶段包含一个下采样层（除了第一个阶段），然后是一个 VSS（视觉状态空间）块的堆叠。VSS 块作为表示学习的视觉对应物，类似于 Mamba 块 [17]（见图 3 (b)）。VSS 块的初始架构（在图 3 (c) 中称为"原始 VSS 块"）是通过替换 S6 模块来构建的。S6 是 Mamba 的核心，能够实现全局感受野、动态权重（即选择性）和线性复杂度。我们用新提出的二维选择性扫描（SS2D）模块替换它，更多细节将在下面小节中介绍。为了进一步提高计算效率，我们去除了整个乘法分支（在图 3 (c) 中用红色框突出显示），因为门控机制的效果已经通过 SS2D 的选择性实现。因此，改进后的 VSS 块（见图 3 (d)）由一个网络分支和两个残差模块组成，模仿了原始 Transformer 块 [58] 的架构。本文中的所有结果均使用基于这种架构构建的 VSS 块的 VMamba 模型获得。

1. 输入处理

输入图像 ：输入图像是一个包含高度 H、宽度 W 和三个颜色通道（RGB）的三维张量。

划分小块 ：图像首先通过一个 stem module 被划分为多个小块（patches），生成一个空间维度为的二维特征图。这种划分方式使得模型能够处理图像的局部特征。

2. 特征提取阶段

多阶段网络 ：模型采用多个网络阶段，创建不同分辨率的分层表示。各个阶段的分辨率依次为 ,和。

下采样层：每个阶段（除了第一个阶段）都包含一个下采样层，用于逐步减小特征图的尺寸，从而捕获更高层次的特征。

3. 视觉状态空间块（VSS Block）

VSS 块结构：每个阶段的核心是一个或多个视觉状态空间（VSS）块。VSS 块的设计灵感来源于 Mamba 块(图3(b)) [17]，用于进行高效的表示学习。

原始 VSS 块 ：最初的 VSS 块（图3(c)被称为"原始 VSS 块"）通过替换 S6 模块来构建，该模块是 Mamba 的核心，能够实现全局感受野、动态权重和线性复杂度。

4. 2D-选择性扫描（SS2D）模块

动态选择性 ：为了适应视觉数据的特性，VSS 块中的 S6 模块被替换为新的 SS2D 模块(图3(d))，后者能够处理二维数据并保留选择性扫描的优势。

数据转发过程：SS2D 模块通过三个步骤进行数据转发：交叉扫描（Cross-Scan）、选择性扫描（Selective Scanning）和交叉合并（Cross-Merge），使得图像中每个像素都能整合来自不同方向的信息，以建立全局感受野。

5. 结构优化

简化架构：为了提高计算效率，去除了整个乘法分支，因为 SS2D 的选择性机制已经实现了门控机制的效果。改进后的 VSS 块由一个单一的网络分支和两个残差模块组成，模仿了原始 Transformer 块的架构。

层配置：通过优化层的配置（例如减少深度可分离卷积层的数量），VMamba 实现了更高的计算效率，并降低了 FLOPs（每秒浮点计算次数）。

4.2 视觉数据的二维选择性扫描（SS2D）

虽然 S6 中的扫描操作的顺序特性与涉及时间数据的 NLP 任务相匹配，但将其应用于视觉数据时却带来了显著挑战，因为++视觉数据本质上是非顺序的，包含空间信息（例如局部纹理和全局结构）++。

为了解决这个问题，++S4ND [44] 通过卷积操作重新构造 SSM，直接将核从 1D 扩展到 2D，采用外积的方法。然而，这种修改限制了权重的输入依赖性，导致捕获上下文信息的能力有限。++

因此，我们坚持采用选择性扫描方法 [17] 进行输入处理，并提出 SS2D 模块，以在不妥协其优势的前提下将 S6 适应于视觉数据。

交叉扫描（Cross-Scan）

功能：

在此步骤中，输入的特征块（patches）被展开成多个序列。这个展开是沿着四个不同的遍历路径进行的，这样每个块能够从不同的方向获得信息。

实现方式：

将特征块按行和列进行遍历，形成四条扫描路径。例如，可能的路径包括从左到右、从右到左、从上到下和从下到上。

这一过程确保了每个小块都有机会与邻近块进行交互，从而增强了特征的上下文理解能力。

选择性扫描（Selective Scanning）

核心组件：

在此步骤中，展开后的序列会通过多个 S6 块 进行处理。每个 S6 块独立处理其对应的序列，从而实时提取特征。

动态权重：

S6 块能够根据输入数据的特性调整权重，实现动态选择性。这使得模型能够聚焦于输入中最相关的特征，从而提高表示的有效性。

并行处理：

每个 S6 块的并行处理允许高效的信息提取，加快了特征处理的速度，减少了计算瓶颈。

交叉合并（Cross-Merge）

合并步骤：

在选择性扫描之后，生成的序列会被重新整形并合并，形成最终的输出特征图。这个步骤的目的是将从不同路径获取的信息整合到一起。

输出特征图：

通过交叉合并，输出的特征图不仅保留了各个块的局部特征，还整合了全局信息，使得每个像素能够利用来自各个方向的上下文信息。

图 2 表示 SS2D 中的数据传递分为三个步骤：交叉扫描（Cross-Scan） 、使用S6 块的选择性扫描 和交叉合并（Cross-Merge）。具体来说，SS2D 首先沿四个不同的遍历路径将输入块展开为序列（即交叉扫描）。然后，每个块序列都使用单独的 S6 块并行处理，结果序列被重新调整形状并合并形成输出特征图（即交叉合并）。通过使用互补的一维遍历路径，SS2D 使图像中的每个像素能够从不同方向整合来自其他像素的信息。这种集成促进了在二维空间中建立全局感受野。

4.3 加速 VMamba

如图 3 (e) 所示，采用原始 VSS 块的 VMamba-T 模型（称为"原始 VMamba"即"Vanilla VMamba"）实现了每秒处理 426 张图像，参数量为 22.9M，FLOPs 为 5.6G。尽管在 Tiny 级别上达到了 82.2% 的前沿分类准确率（比 Swin-T [36] 高 0.9%），但较低的吞吐量和较高的内存开销对 VMamba 的实际部署构成了重大挑战。

在这一小节中，我们概述了提升推理速度的努力，主要集中在实现细节和架构设计的改进。我们在 ImageNet-1K 上评估模型的图像分类。每项逐步改进的影响总结如下，其中（%，img/s）表示在 ImageNet-1K 上的 Top-1 准确率提升和推理吞吐量的提升。附录 E 中提供了进一步讨论。

步骤 (a) (+0.0%, +41 img/s)：通过在 Triton 中重新实现交叉扫描和交叉合并。

步骤 (b) (+0.0%, −3 img/s)：调整选择性扫描的 CUDA 实现，以适应 float16 输入和 float32 输出。这显著提高了训练效率（吞吐量从 165 提升至 184），尽管测试时速度略有波动。

步骤 (c) (+0.0%, +174 img/s)：通过将选择性扫描中的相对较慢的 einsum 替换为线性变换（即 torch.nn.functional.linear）。我们还采用了 (B, C, H, W) 的张量布局，以消除不必要的数据排列。

步骤 (d) (−0.6%, +175 img/s)：通过引入 MLP 到 VMamba 中，因其计算效率。我们还去掉了 DWConv（深度可分离卷积 [23]）层，并将层配置从 [2,2,9,2] 改为 [2,2,2,2] 以降低 FLOPs。

步骤 (e) (+0.6%, +366 img/s)：将参数 ssm-ratio（特征扩展因子）从 2.0 降低到 1.0（也称为步骤 (d.1)），将层数提高到 [2,2,5,2]（也称为步骤 (d.2)），并去掉整个乘法分支，如图 3 (c) 所示。

步骤 (f) (+0.3%, +161 img/s)：通过引入 DWConv 层（也称为步骤 (e.1)）并将参数 d_state（SSM 状态维度）从 16.0 降低到 1.0（也称为步骤 (e.2)），同时将 ssm-ratio 重新提高到 2.0。

步骤 (g) (+0.1%, +346 img/s)：通过将 ssm-ratio 降低到 1.0，同时将层配置从 [2,2,5,2] 更改为 [2,2,8,2]。

五、实验

在本节中，我们呈现了一系列实验，以评估 VMamba 的性能，并将其与各种视觉任务中的流行基准模型进行比较。我们还验证了提出的二维特征图遍历方法的有效性，比较了替代方法。此外，我们通过可视化有效感受野（ERF）和激活图，分析了 VMamba 的特性，并考察了其在较长输入序列上的可扩展性。我们主要遵循 Swin [36] 中使用的超参数设置和实验配置。详细的实验设置请参见附录 E 和 F，其他消融实验见附录 H。++所有实验均在配备 8 个 NVIDIA Tesla-A100 GPU 的服务器上进行。++

5.1 图像分类

我们在 ImageNet-1K [9] 上评估了 VMamba 的图像分类性能，比较结果汇总在表 1 中。

VMamba-T 在相似 FLOPs 的情况下，实现了 82.6% 的 top-1 准确率，超越了 DeiT-S 2.8% 和 Swin-T 1.3%。值得注意的是，VMamba 在 Small 和 Base 规模上都保持了性能优势。例如，VMamba-B 实现了 83.9% 的 top-1 准确率，超越了 DeiT-B 2.1% 和 Swin-B 0.4%。在计算效率方面，VMamba-T 达到 1,686 张图像/秒的吞吐量，优于或与最先进的方法相当。VMamba-S 和 VMamba-B 的吞吐量分别为 877 张图像/秒和 646 张图像/秒。与基于 SSM 的模型相比，VMamba-T 的吞吐量比 S4ND-Conv-T [44] 高 1.47 倍，比 Vim-S [69] 高 1.08 倍，同时分别保持了 0.4% 和 2.1% 的明显性能领先。

5.2 下游任务

在这一小节中，我们评估了 VMamba 在下游任务上的性能，包括 MSCOCO2017 [33] 上的目标检测和实例分割，以及 ADE20K [68] 上的语义分割。训练框架基于 MMDetection [3] 和 MMSegmentation [4] 库，遵循 [35] 使用 Mask R-CNN [26] 和 UperNet [63] 作为检测和分割网络。

目标检测和实例分割

在 MSCOCO 上的结果呈现在表 2 中。

VMamba 在不同训练周期中，在箱体和掩膜平均精度（和）上表现出色。在 12 轮微调的情况下，VMamba-T/S/B 分别达到了 47.3%/48.7%/49.2% 的目标检测 mAP，超越了 Swin-T/S/B 分别 4.6%/3.9%/2.3% 的 mAP，以及 ConvNeXt-T/S/B 分别 3.1%/3.3%/2.2% 的 mAP。此外，VMamba-T/S/B 在实例分割 mAP 上也超过了 Swin-T/S/B 分别 3.4%/2.8%/1.8% 的 mAP，以及 ConvNeXt-T/S/B 分别 2.6%/1.9%/1.4% 的 mAP。此外，VMamba 的优势在使用多尺度训练的 36 轮微调周期中持续存在，突显其在需要密集预测的下游任务中的强大潜力。

语义分割

与之前的实验一致，VMamba 在 ADE20K 上的语义分割性能优于其他模型，同时参数量相当。如表 2 所示，在单尺度（SS）设置下，VMamba-T 的 mIoU 比 Swin-T 高 3.4%，比 ConvNeXt-T 高 1.9%，并且在多尺度（MS）输入下优势持续存在。对于 Small 和 Base 级别的模型，VMamba-S/B 在 SS 设置下分别比 NAT-S/B [25] 高 2.6%/2.5% 的 mIoU，在 MS 设置下分别高 1.7%/1.9% 的 mIoU。

讨论

这一小节的实验结果展示了 VMamba 在目标检测、实例分割和语义分割任务中的适应能力。

在图 4 (a) 中，我们将 VMamba 的性能与 Swin 和 ConvNeXt 进行比较，强调其在处理下游任务时的优势，同时在 ImageNet-1K 上保持相似的分类准确率。这一结果与图 4 (b) 一致，VMamba 在不同输入图像尺寸下表现出最稳定的性能（即，性能下降幅度适中），在 768 × 768 的输入分辨率下，无需微调时实现了 74.7% 的 top-1 分类准确率（线性调优后为 79.2%）。虽然对输入分辨率变化的容忍度更高，VMamba 仍保持 FLOPs 和内存消耗的线性增长（见图 5 (a) 和 (c)），并保持高吞吐量（见图 5 (b)），使其在适应具有更大空间分辨率的下游任务时，相较于基于 ViT 的方法更为有效和高效。这与 Mamba 在高效长序列建模方面的先进能力 [17] 一致。

5.3 分析

SS2D 与自注意力的关系

在上述公式 5 中，涉及 Q、K和 V 的矩阵乘法过程，与自注意力机制非常相似，尽管包括了额外的权重 w。

激活图的可视化

为了更直观和深入地理解 SS2D，我们进一步可视化特定查询块（称为"激活图"）对应的的注意值。

如图 6 (b) 所示，的激活图展示了 SS2D 在捕获和保留遍历信息方面的有效性，所有之前扫描的标记在前景区域都被激活。此外，w 的引入导致激活图更加集中在查询块的邻域（图 6 (c)），这与 w 的公式中固有的时间加权效应一致。然而，选择性扫描机制允许 VMamba 沿着扫描路径累积历史，从而促进跨图像块的长期依赖关系的建立。这在红框圈起的子图中显现出来（图 6 (d)），++左侧远处的羊块（在更早步骤中扫描）仍然保持激活。++

注意，这里是做了两次实验。第一次实验使用左边的羊作为查询块，第二次使用右边的羊作为查询块。每一行结果针对一次实验（一个查询块位置）生成的激活图。

(a) 输入图像

图像显示了两只羊，红色五角星表示模型设定的查询块（Query Patch），即模型特别关注的输入区域。

查询块是模型分析和生成激活图的起始点。

(b) 激活图

含义：该激活图表示 SS2D 中的 Self-Attention 机制，仅基于标准的****（查询和键之间的相似性计算）。

特点：

激活区域分布较广，覆盖了图像中多个区域，尤其是前景（羊）部分。

该激活图有效展示了模型捕获的全局信息，但由于缺少额外的选择机制，可能存在噪声。

(c) 激活图

含义：这是 SS2D 中包含权重 w 的选择性扫描机制的激活图。

特点：

激活区域更集中于查询块周围的邻域，反映了 w 对局部区域的加权效果。

与 (b) 相比，激活图更具选择性和聚焦性，噪声减少。

这种机制通过时间加权作用，使模型能够更有效地捕获重要信息。

(d) 对单一路径的激活图

含义：这一部分展示了在单一扫描路径（例如从左到右、从上到下等）上的激活图。

特点：

每列对应一种扫描路径（图中用红色箭头标注扫描方向）：左到右、右到左、上到下、下到上。

每种扫描路径都针对不同方向捕获了特定的信息。

红色框部分可视化目的是说明 选择性扫描机制 如何逐步沿扫描路径累积信息，并能够激活左侧的羊（即较早扫描的区域）。这表明模型在进行扫描时不会遗忘早期的上下文信息，从而有效建立了 长距离依赖。

有关更多可视化和进一步讨论，请参见附录 D。

有效感受野的可视化

有效感受野（ERF） [40, 11] 是指对特定输出单元激活贡献的输入空间区域。我们对不同视觉骨干网络的中心像素 ERF 进行了比较分析，比较训练前后的结果呈现在图 7 中。

颜色表示激活强度：

图中由浅绿色到深绿色（颜色越深越接近 0.8）表示中心像素对输入图像不同区域的响应强度。

深绿色（高强度）表示该区域对中心像素激活有较大贡献。

浅绿色（低强度）表示贡献较低。

接近白色（零强度**）** 表示该区域对中心像素的激活几乎没有影响。

区域表示感受野的覆盖范围：

每个模型的 ERF 中绿色区域的大小代表中心像素可以感知到的输入图像区域。

绿色区域越大，说明模型的 ERF 覆盖能力越强，能捕获更多的上下文信息。

绿色区域越小，说明模型的 ERF 更局限于局部特征，覆盖能力较差。

在所检查的模型中，只有 DeiT、HiViT、Vim 和 VMamba 显示出了全局 ERF，而其他模型尽管理论上具有全局覆盖的潜力，但却表现出局部 ERF。

DeiT、HiViT、Vim 和 VMamba这些模型的绿色区域覆盖了整个输入图像，即它们能捕获输入图像的全局信息。

ResNet-50 或 ConvNeXt 的 ERF（训练前后）：这些模型的绿色区域集中在图像中心附近，说明它们的 ERF 是局部的，主要依赖局部信息进行决策。

此外，VMamba 的线性时间复杂度使其计算效率优于 DeiT 和 HiViT，后者在输入块数量上呈现二次成本。

DeiT 和 HiViT：由于它们基于全局注意力机制，）处理更多输入块，覆盖整个图像的 ERF，计算复杂度较高（O(n2)），处理高分辨率图像或更多输入块时效率较低。

VMamba： 通过设计优化（如 SS2D 和局部加权机制），实现了 线性时间复杂度（O(n)）（训练前输入局部图像，部分绿色，在保持全局 ERF 的同时，计算效率显著提升。（

虽然 VMamba 和 Vim 都基于 Mamba 架构，但 VMamba 的 ERF 更均匀并且具有 2D 识别能力，这可能直观地解释了其优越的性能。

VMamba 的 ERF 颜色分布更加平滑且均匀（绿色区域的颜色深浅逐步过渡），覆盖了整个图像区域。表明 VMamba 能够在整个图像中均衡地捕捉特征，而不是仅仅聚焦在某些局部区域。

VMamba 的 ERF 显示为一个规则且对称的 2D 分布形状，即圆形或接近均匀扩散的分布模式。这种规则的 2D 显示表明 VMamba 能够在水平方向和垂直方向上均衡地捕获特征，而不是只在单一方向上具有强感知能力。

选择性扫描模式的诊断研究

我们将提出的扫描模式（即交叉扫描 ）与三种基准模式进行比较：单向扫描 （Unidi-Scan）、双向扫描 （Bidi-Scan）和级联扫描（Cascade-Scan，依次横向和纵向扫描数据）。

特征维度经过调整，以保持相似的架构参数和 FLOPs 以保证公平比较。

蓝色表示有 DWConv 层，红色表示去除 DWConv 层。

如图 8 所示，交叉扫描在++计算效率++ 和++分类准确率++方面均优于其他扫描模式，突显其在实现 2D 选择性扫描方面的有效性。去掉 DWConv 层（已被证明有助于模型学习 2D 空间信息）进一步增强了这一优势。这强调了交叉扫描通过采用四向扫描捕获 2D 上下文信息的内在优势。

六、结论

本文介绍了 VMamba，这是一种高效的视觉骨干模型，基于状态空间模型（SSMs）构建。VMamba 将来自自然语言处理任务的选择性 SSMs 的优势整合到视觉数据处理中，弥合了有序的一维扫描与非顺序二维遍历之间的差距，采用了新颖的 SS2D 模块。此外，我们通过一系列架构和实现的改进显著提高了 VMamba 的推理速度。通过大量实验，已经证明了 VMamba 系列的有效性，其线性时间复杂度使得 VMamba 在处理大分辨率输入的下游任务中具有优势。

限制

尽管 VMamba 展现出良好的实验结果，但仍有改进的空间。先前的研究已经验证了在大规模数据集（如 ImageNet-21K）上无监督预训练的有效性。然而，现有的预训练方法与基于 SSM 的架构（如 VMamba）的兼容性，以及为这些模型量身定制的预训练技术的识别，仍然未被探索。研究这些方面可能为架构设计的未来研究开辟了有希望的方向。此外，有限的计算资源限制了我们对 VMamba 架构在大规模下的探索，以及进行了细致的超参数搜索以进一步提升实验性能。尽管 VMamba 的核心组件 SS2D 并未对输入数据的布局或模态做出特定假设，允许其在各种任务中进行推广，但 VMamba 在更广泛任务中的集成潜力仍未被探索。弥合 SS2D 与这些任务之间的差距，以及为视觉任务提出更通用的扫描模式，代表了一个有前景的研究方向。

附录内容待补充......

Wait~