VMamba: 视觉状态空间模型.
0. TL;DR
本文提出了VMamba，一种用于视觉表示学习的视觉状态空间模型。VMamba通过引入全局感受野和动态权重，实现了线性复杂度，同时保持了视觉建模的高性能。通过引入交叉扫描模块（CSM），VMamba解决了方向敏感性问题，使其能够有效地处理视觉图像。实验结果表明，VMamba在图像分类、目标检测和语义分割等任务上均优于现有的CNN和视觉Transformer模型，并且在图像分辨率增加时表现出更显著的优势。
1. 背景介绍
视觉表示学习是计算机视觉中的一个基本研究课题。自深度学习时代以来，卷积神经网络（CNNs）和视觉Transformer（ViTs）一直是两种最主流的基础模型。CNNs在图像分辨率方面表现出线性复杂度，而ViTs则在拟合能力上优于CNNs，尽管它们面临二次复杂度的挑战。ViTs通过引入全局感受野和动态权重，实现了更优越的视觉建模性能。然而，这种全局感受野和动态权重的引入也带来了计算复杂度的增加，尤其是在处理高分辨率图像时。
为了在保持全局感受野和动态权重的同时提高计算效率，本文提出了一种基于状态空间模型（SSMs）的新型架构，即视觉状态空间模型（VMamba）。VMamba通过引入选择性扫描机制，实现了线性复杂度，同时保持了全局感受野。此外，为了处理非因果视觉图像，VMamba引入了交叉扫描模块（CSM），使其能够有效地捕获图像中的全局信息。
2. VMamba 模型
VMamba是一种集成了基于Mamba模块的视觉骨干网络，以促进高效的视觉表示学习。Mamba的并行选择扫描操作是为处理一维顺序数据而设计的。当尝试将其适用于视觉数据处理时会有问题，因为视觉数据本质上缺乏视觉组件的顺序排列。
为了解决这个问题，作者提出了二维选择扫描（SS2D），这是一种为空间域遍历量身定制的四向扫描机制，促进选择性SSM在处理视觉数据方面的扩展。与自注意力机制相比，SS2D确保每个图像块仅通过沿相应扫描路径计算的压缩隐藏状态获取上下文知识，从而将计算复杂度从二次降低到线性。
VMamba 的架构概览如图所示：
- 输入图像$\mathbf{I} \in \mathbb{R}^{H \times W \times 3}$首先通过一个patch partition模块划分成图像块，生成一个空间维度为$H/4 × W/4$的二维特征图；
- 多个阶段被用来创建分辨率为$H/8 × W/8$、$H/16 × W/16$ 和$H/32 × W/32$ 的分层表示，每个阶段包括一个下采样层和一个视觉状态空间（VSS）模块。
- VSS模块设计受Mamba模块的启发，通过引入SS2D扫描机制来处理二维数据；
- 为了进一步提高计算效率，作者消除了整个乘法分支，最终VSS模块由一个具有两个残差模块的单一网络分支组成。
SS2D中的数据转发涉及三个步骤：
- 交叉扫描：给定输入数据，SS2D首先沿四条不同的遍历路径（即交叉扫描，分别从二维矩阵的左上角向右$\rightarrow$、左上角向下$↓$、右下角向左$←$、右下角向上$↑$出发）；
- 选择性扫描：使用单独的S6模块并行处理每个序列；
- 交叉合并：随后重新整形并合并结果序列以形成输出图。
通过采用互补的1D遍历路径，SS2D使图像中的每个像素能够有效整合来自其他像素的不同方向的信息，从而促进在2D空间中建立全局感受野。
使用原始VSS模块的VMamba-T模型实现了426张图像/秒的吞吐量，包含22.9M参数和5.6G FLOPs。尽管在微型级别上实现了82.2%的最新分类准确率，但低吞吐量和高内存开销对VMamba的实际部署构成了重大挑战。作者通过在ImageNet-1K上的图像分类对模型进行评估。测试了一系列渐进改进：
- step (a) (+0.0%,+41img/s)：通过在Triton中重新实现Cross-Scan和Cross-Merge以适应float16输入和float32输出；
- step (b) (+0.0%，−3img/s)：通过用线性变换（
torch.nn.functional.linear）替换选择性扫描中相对较慢的einsum，尽管测试时间略有速度波动，但这显著提高了训练效率（吞吐量从165提升到184）； - step (c) (+0.0%，+174img/s)：采用$(B, C, H, W)$的张量布局以消除不必要的尺寸转换；
- step (d) (−0.6%, +175img/s)：引入MLP到VMamba中，舍弃了DWConv层，并将层配置从$[2,2,9,2]$更改为$[2,2,2,2]$以降低FLOPs；
- step (e) (+0.6%, +366img/s)：舍弃整个乘法分支，并将参数sm-ratio（特征扩展因子）从2.0降低到1.0，使得可以将层数增加到$[2,2,5,2]$，同时减少FLOPs；
- step (f) (+0.3%, +161 img/s)：将参数 d_state（SSM 状态维度）从 16.0 降低到 1.0，使得可以将 ssm-ratio 提回到 2.0，并在不增加 FLOPs 的情况下引入 DWConv 层；
- step (g) (+0.1%, +346 img/s)：将 ssm-ratio 降低到 1.0，同时将层配置从 $[2,2,5,2]$ 更改为$[2,2,8,2]$。
3. 实验分析
对于图像分类任务，VMamba在ImageNet-1K数据集上的表现优于现有的CNN和视觉Transformer模型。具体来说，VMamba-T在224×224输入尺寸下达到了82.6%的Top-1准确率，超过了RegNetY-4G、DeiT-S和Swin-T等模型。
对于目标检测任务，VMamba在COCO数据集上的表现同样优于现有模型。在12轮和36轮训练中，VMamba-T、VMamba-S和VMamba-B模型在目标检测和实例分割任务上均取得了更高的AP值。
对于语义分割任务，VMamba在ADE20K数据集上的表现同样优于现有模型。在512×512和640×640输入尺寸下，VMamba-T、VMamba-S和VMamba-B模型在语义分割任务上均取得了更高的mIoU值。
VMamba的有效感受野（ERF）显示出其全局特性，这与其他模型如ResNet、ConvNeXt、Swin和DeiT进行了对比。VMamba在训练后展现出全局ERF，表明其对图像的全局感知能力。
VMamba在不同输入尺寸下的性能表现出更稳定的趋势，随着输入尺寸的增加，VMamba的性能逐渐提升，显示了其对输入图像尺寸变化的鲁棒性。
