论文标题：VMamba: Visual State Space Model
论文地址：https://arxiv.org/pdf/2401.10166

摘要
VMamba 是一个视觉骨干网络，基于状态空间模型（SSM），其复杂度是线性的。该架构的核心是视觉状态空间（VSS）模块和 2D 选择性扫描（SS2D）模块。通过四条扫描路径的遍历，SS2D 实现了从不同方向获取上下文信息的能力，同时降低了计算成本。
论文的贡献主要包括：

• VMamba 作为一种新型的 SSM 基础视觉网络，以线性时间复杂度处理视觉表示学习任务。
  提出了 2D 选择性扫描（SS2D），实现了从1D数组扫描向2D平面遍历的扩展。
• VMamba 在图像分类、目标检测和语义分割任务中表现出色，且在输入尺度扩展方面表现出优势。

架构

VMamba 的架构依赖于 Visual State Space (VSS) Block 模块，通过 SS2D 模块来完成高效的图像表示学习。下图描述了 VMamba 的基本网络结构，其核心思想是将输入图像分割成补丁，并通过多个下采样阶段和 VSS 模块逐步提取层级特征。

主要步骤：

1. 将输入的图像 H x W x 3分成块（图中stem步骤，与patch partition相同），生成空间尺寸为H/4 x H/4 x C_1 的特征图。
2. 层及特征提取：多个网络阶段用来生成分辨率逐步降低的低级特征表示，特征分辨率为 H/8 x W/8，H/16 x W/16, H/32 x W/32。这些阶段，除了第一个阶段，每个阶段包含下采样层，并有多个VSS模块。
3. VSS Block： VSS模块被设计为视觉数据中对应Mamba模块的部分，用来学习图像的特征表示。在VSS模块中，SS2D模块被用来替换S6模块。SS2D可以扫描不同方向的信息来有效提取上下文特征，并且维持线性计算复杂度。

为了提高计算效率，VSS模块中的整个乘法部分被移除掉了，因为选择性机制已经通过SS2D实现了。最终的VSS模块包含一个单一的网络分支，采用两个残差模块，类似Transformer的基本块结构。

SS2D (2D-Selective-Scan)

1. 交叉扫描 (Cross-Scan)
   输入的图像被划分成多个非重叠的补丁（patch），然后沿着四个不同的扫描路径展开（如图所示的红色箭头），每个扫描路径形成一个序列，将图像的 2D 数据转化为一维的序列进行处理。

2. 选择性扫描 (Selective Scanning with S6 Blocks)
   每条扫描路径经过独立的 S6 Block 进行特征提取。S6 模块负责处理这些展开的序列，通过应用状态空间模型（SSM）的公式（如图中右侧部分展示的公式）进行特征提取。该过程允许在不同的方向上整合信息，保留重要的上下文信息并滤除不相关的内容，从而在保持线性复杂度的同时实现全局感受野。

3. 交叉合并 (Cross-Merge)
   选择性扫描后，四个扫描方向的序列被重新组合，通过逆向操作将一维序列还原成二维图像补丁。最终输出的补丁包含了从不同方向整合的上下文信息，完成全局上下文的整合。

VMamba 模型的加速优化

VMamba 是一种用于图像分类的视觉模型。这个模型在未经优化的情况下，尽管精度很高（82.2%），但其推理速度（吞吐量 426 张图像/秒）较慢，同时计算复杂度（5.6 GFLOPs）较高。这些瓶颈会影响它在实际部署中的效率。因此，研究人员对模型进行了逐步优化，以提高其速度和减少计算负担，同时尽量保持较高的分类精度。

优化过程详解：
实现层面的优化：

Step (a) - CSM in Triton:

目标: 改进 Cross-Scan 和 Cross-Merge 操作的实现。
效果: 推理速度提高了 41 张图像/秒，但分类精度不变。通过更高效的内存处理，提高了数据处理速度。
Step (b) - f16 in & f32 out:

目标: 改用 float16 数据格式作为输入，并以 float32 输出，旨在提高训练效率。
效果: 训练时加速效果显著，但测试时速度略有波动（-3 张图像/秒）。精度仍保持不变。
Step © - Einsum → Linear Tensor Layout:

目标: 将模型中慢速的 einsum 操作替换为更高效的线性变换（linear transformation）。同时调整张量布局，减少数据重排操作。
效果: 推理速度提高了 174 张图像/秒，精度保持不变。这大幅加快了推理速度。
架构层面的优化：

Step (d) - MLP and Fewer Layers:

目标: 引入多层感知器（MLP）以取代部分复杂操作，减少层数以降低 FLOPs。
效果: 尽管推理速度提升了 175 张图像/秒，但分类精度略微下降了 0.6%。
Step (e) - No Skip Branch, More Layers:

目标: 删除模型中的乘法分支，并提高模型的层数配置，同时减少 FLOPs。
效果: 速度显著提高了 366 张图像/秒，精度反而提高了 0.6%，这是一个重要的改进。
Step (f) - Reduce d_state, Reintroduce DWConv:

目标: 将 d_state（SSM 状态维度）从 16 降低到 1，并重新引入深度卷积层（DWConv）。
效果: 推理速度提高了 161 张图像/秒，分类精度增加了 0.3%。
Step (g) - Further Reduce ssm-ratio, More Layers:

目标: 进一步降低 ssm-ratio，并增加模型层数配置。
效果: 推理速度再次提高了 346 张图像/秒，精度提高了 0.1%。

实验评估

图像分类

数据集: 使用了 ImageNet-1K 数据集。
性能比较: 在相似的 FLOPs 下，VMamba-T 的 top-1 准确率达到了 82.6%，比 DeiT-S 高 2.8%，比 Swin-T 高 1.3%。
VMamba 在 Small 和 Base 级别上也保持了性能优势。VMamba-B 的 top-1 准确率为 83.9%，比 DeiT-B 高 2.1%，比 Swin-B 高 0.4%。
推理效率: VMamba-T 的吞吐量为 1,686 张图像/秒，性能远超 SSM 模型，如比 S4ND-Conv-T 高 1.47 倍。

下游任务

物体检测与实例分割:
数据集: MSCOCO2017 数据集。
方法: 使用 Mask-RCNN 和 UperNet 进行物体检测和分割。
结果:
在 12 轮 fine-tuning 训练中，VMamba-T/S/B 的物体检测 mAP 达到 47.3%/48.7%/49.2%，分别比 Swin-T/S/B 高 4.6%/3.9%/2.3% mAP。
实例分割任务中，VMamba 的 mAP 同样超越 Swin 和 ConvNeXt。
在 36 轮的 fine-tuning 训练中，VMamba 继续表现出色，显示了其在密集预测任务中的潜力。
语义分割:

数据集: ADE20K 数据集。
结果:
在单尺度设置下，VMamba-T 的 mIoU 比 Swin-T 高 3.4%，比 ConvNeXt-T 高 1.9%。
在小模型和大模型级别上，VMamba-S/B 在多个设置下也优于其他方法。

结果分析:

VMamba 展示了其在物体检测、实例分割和语义分割中的适应性。
相较于 ViT 模型，VMamba 在处理更大空间分辨率的输入时，表现出 线性增长的 FLOPs 和较小的准确率下降，使其在处理大输入时更加高效。
通过这些实验，VMamba 展现了其在多任务中的强大适应能力，尤其是在处理高分辨率输入和密集预测任务时，具有很好的表现和效率。