VMamba: Visual State Space Model 论文总结

题目：VMamba: Visual State Space Model（视觉状态空间模型）

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

源码：https://arxiv.org/pdf/2401.10166 (github.com)

一、摘要

二、引言

三、方法

3.1. Network Architecture

3.2. 2D-Selective-Scan for Vision Data (SS2D，2D 选择扫描)

3.3. Accelerating VMamba（加速 VMamba）

四、实验

一、摘要

研究背景：在计算机视觉领域，设计计算效率高的网络架构一直是一种需要。

主要工作：本文将状态空间语言模型 Mamba 移植到 VMamba 中，VMamba 是一种线性时间复杂度的视觉骨干。

VMamba 的核心是 视觉状态空间(VSS)模块 的堆叠和 2D选择性扫描 (SS2D) 模块。
通过沿4条扫描路线遍历，SS2D 有助于弥合 1D 选择性扫描的有序性质与 2D 视觉数据的非顺序结构之间的差距，有利于从不同来源和角度收集上下文信息。
基于VSS模块，开发了一系列VMamba架构，并通过一系列架构和实现的增强来加速它们。

实验效果：广泛的实验展示了VMamba在各种视觉感知任务中的良好性能，突出了其与现有基准模型相比在输入缩放效率方面的优势。

二、引言

ViT概述（优缺点）： ViT通常在大规模数据上表现出上级的学习能力。然而，自注意的二次复杂性导致受限。在涉及大空间分辨率的下游任务中会引入了相当大的计算开销。

研究现状：为了应对这一挑战，许多方法已经做出了相当大的努力来提高注意力计算的效率。然而，现有的方法要么对有效感受野的大小施加限制，要么在不同的任务中经历明显的性能下降。

Mamba概述：最近，在自然语言处理(NLP)领域，Mamba 是一种新的状态空间模型(SSM)，已经成为具有线性复杂度的长序列建模的一种非常有前途的方法。

研究问题：然而，Mamba的核心算法 —— 并行化选择扫描操作本质上是为处理一维序列数据而设计的。当试图将其用于处理视觉数据时，这提出了一个挑战，因为视觉数据天生缺乏视觉组件的顺序排列。

主要工作：为了解决这个问题，本文提出了一种 面向空间域遍历的四路扫描机制 SS2D (2D Selective Scan) 。与自注意力机制相比，SS2D 确保每个图像 patch 通过沿着相应的扫描路径计算的压缩隐藏状态独占地获得上下文知识，从而将计算复杂度从二次降低到线性。

实验效果：具体而言，在 ImageNet-1K 上，VMamba-Base 的准确率达到了83.9%，排在第一位，超过 Swin +0.4%，吞吐量大大超过 Swin 40%（646对458）。Vmamba 的优越性扩展到各种下游任务，其中 Vmamba-Tiny/Small/Base 在 COCO 上的目标检测中实现了 47.3%/48.7%/49.2% 的 mAP（1×训练计划）。这分别超过 Swin 4.6%/3.9%/2.3% 和 ConvNeXt 3.1%/3.3%/2.2% 的表现。对于 ADE20K 上的单尺度语义分割，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 在保持相当性能的同时，FLOP 呈线性增长。

主要贡献：

提出了 VMamba，一个基于SSM的视觉骨干网络，用于具有 线性时间复杂度 的视觉表示学习。为了提高VMamba的推理速度，在体系结构设计和实现细节上进行了一系列改进。
引入了 二维选择性扫描（SS2D）来弥补一维阵列扫描和二维平面遍历之间的差距，便于将选择性扫描扩展到处理视觉数据。
VMamba 在一系列视觉任务中表现出了良好的性能，包括图像分类、对象检测和语义分割。它还表现出对输入序列长度的显著适应性，显示出计算复杂度的线性增长。

三、方法

3.1. Network Architecture

流程：输入图像 $\mathrm{I \in R^{H \times W \times 3}}$ 首先被 stem模块分割成 patch，形成空间维度为 H/4 × W/4 的 2D 特征图。随后，采用多个网络阶段来创建 H/8 × W/8、H/16 × W/16 和 H/32 × W/32 分辨率的分层表示。每个阶段包括一个下采样层 (除了第一阶段)，然后是视觉状态空间(VSS)模块的堆叠。（Stem模块 + 下采样层 + VSS模块）

VSS模块：为了进一步提高计算效率，消除了整个乘法分支 (图中红色框)，因为门控机制的效果是通过SS2D的选择性实现的。因此，得到的 VSS 模块 (如图(d)所示) 由带有两个残差模块的单个网络分支组成，模仿了普通Transformer模块的体系结构。

3.2. 2D-Selective-Scan for Vision Data (SS2D，2D 选择扫描)

问题：虽然 S6 中的扫描操作的顺序性质与涉及时间数据的 NLP 任务很好地一致，但当应用于视觉数据时，其提出了显著的挑战，视觉数据本质上是非顺序的并且包含空间信息（例如，局部纹理和全局结构）。

解决方法：坚持选择性扫描方法进行输入处理，并提出了 2D 选择性扫描（SS2D）模块，以使 S6 适应视觉数据，而不损害其优势。

过程：给定输入数据，SS2D 首先 沿着四个不同的遍历路径 将输入 patch 展开为序列（即，交叉扫描），使用单独的S6模块并行地处理每个 patch 序列，并随后对所得到的序列进行整形和合并以形成输出图（即，交叉合并）。通过采用互补的1D遍历路径，SS2D 使得图像中的每个像素能够有效地整合来自不同方向上的所有其他像素的信息，从而有助于在2D空间中建立全局感受野。

如下图所示，SS2D 中的 data forwarding （前向）包括三个步骤：cross-scan（交叉扫描）、selective scanning（基于S6模块的选择性扫描）和 cross-merge（交叉合并）。

3.3. Accelerating VMamba（加速 VMamba）

通过ImageNet-1 K上的图像分类来评估模型。每个渐进式改进的影响总结如下，其中（%，img/s）分别表示ImageNet-1 K上的 top-1准确度 和 推理吞吐量 的增益。如下所示：

步骤(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.function.linear)。采用 (B, C, H, W) 张量布局来消除不必要的数据排列。
步骤(d) (−0.6%，+175 img/s) 由于其计算效率，将 MLP 引入 VMamba。丢弃了DWConv层，并将层配置从 [2,2,9,2] 更改为 [2,2,2,2]，以降低FLOPs。
步骤(e) (+0.6%， +366 img/s) 丢弃整个乘法支路，并将参数 ssm-ratio(特征扩展因子) 从2.0降低到1.0。这允许层数提高到 [2,2,5,2]，同时减少FLOPs。（证明所提出的VMamba的体系架构的有效性）
步骤(f) (+0.3%， +161 img/s) 通过将 参数 d_state (SSM状态维度) 从16.0减少到1.0。这允许ssm-ratio提高到2.0，并引入DWConv层，而不会增加FLOPs。
步骤(g) (+0.1%， +346 img/s) 将ssm-ratio降低到1.0，同时将层配置从[2,2,5,2]更改为[2,2,8,2]。