VMamba：视觉SSM

论文标题：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 模块逐步提取层级特征。

主要步骤：

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

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

SS2D (2D-Selective-Scan)

交叉扫描 (Cross-Scan)

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

每条扫描路径经过独立的 S6 Block 进行特征提取。S6 模块负责处理这些展开的序列，通过应用状态空间模型（SSM）的公式（如图中右侧部分展示的公式）进行特征提取。该过程允许在不同的方向上整合信息，保留重要的上下文信息并滤除不相关的内容，从而在保持线性复杂度的同时实现全局感受野。
交叉合并 (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 展现了其在多任务中的强大适应能力，尤其是在处理高分辨率输入和密集预测任务时，具有很好的表现和效率。