在计算机视觉领域，设计计算高效的网络架构一直是研究的热点。今天，我想和大家分享一篇发表在 NIPS 2024 上的论文——VMamba：Visual State Space Model，这篇论文提出了一种新的视觉骨干网络，具有线性时间复杂度，展现了在多种视觉感知任务中的出色表现。
源码：GitHub - MzeroMiko/VMamba: VMamba: Visual State Space Models，code is based on mamba
一、亮点解析
（一）2D 选择性扫描（SS2D）模块
作者巧妙地提出了 2D 选择性扫描（SS2D）模块，这一模块通过四向扫描机制，有效地弥合了 1D 选择性扫描与 2D 视觉数据之间的差距。SS2D 确保每个图像块仅通过沿相应扫描路径计算的压缩隐藏状态获得上下文知识，从而将计算复杂度从二次降低到线性。这种创新的设计，犹如在图像的二维世界中开辟了四条 “信息高速通道”，让图像块能够更高效地获取上下文信息，极大地提升了计算效率。
（二）VMamba 架构
基于 VSS 模块，作者精心开发了一系列 VMamba 架构，包括 VMamba-Tiny、VMamba-Small 和 VMamba-Base 等。这些架构通过一系列架构增强和实现优化，实现了加速。VMamba 架构的出现，就像为计算机视觉任务打造了一套 “高性能引擎”，在保持相当性能的同时，计算复杂度呈线性增长，与基于 ViT 的模型在输入 token 数量增加时计算复杂度呈二次增长形成了鲜明对比，为处理大规模视觉数据提供了更高效的选择。下图是只能VMamba整体架构。
二、方法
（一）网络架构的构建
作者开发了三种规模的 VMamba，输入图像首先通过一个 stem 模块分割成 patches，生成一个空间维度为的 2D 特征图。随后，多个网络阶段用于创建分辨率为、和的分层表示。每个阶段（除了第一个阶段）包含一个下采样层，后跟一系列视觉状态空间（VSS）模块。VSS 模块作为 Mamba 模块的视觉对应物，用于表示学习。最终的 VSS 模块由一个包含两个残差模块的单一网络分支组成，模仿了原始 Transformer 模块的架构。
（二）2D 选择性扫描用于视觉数据（SS2D）
视觉数据本质上是非顺序的，包含空间信息。为了解决这一挑战，SS2D 中的数据前向传播涉及三个步骤：交叉扫描、使用 S6 块的选择性扫描和交叉合并。通过采用互补的 1D 遍历路径，SS2D 使图像中的每个像素能够有效地从不同方向整合来自所有其他像素的信息，从而促进在 2D 空间中建立全局感受野。这种设计，仿佛赋予了每个像素 “多角度视野”，让它们能够更全面地感知图像内容。
（三）加速 VMamba 的方法
作者通过在 ImageNet-1K 上的图像分类评估模型，逐步改进 VMamba 的推理速度。例如，通过在 Triton 中重新实现交叉扫描和交叉合并，调整选择性扫描的 CUDA 实现以适应 float16 输入和 float32 输出，将选择性扫描中的相对较慢的 einsum 替换为线性变换等方法，不断提高推理吞吐量。这些优化措施，就像是对 VMamba 进行了一系列 “精细调校”，让它在实际应用中能够更高效地运行。
三、 实验结果
论文评估了视觉状态空间模型 （VMamba） 在各种任务中的作用，包括在ImageNet-1K数据集做图像分类、COCO 数据集做目标检测和在ADE20K 数据集上做语义分割。结果表明，VMamba 实现了强大的性能和效率，性能优于 CNN 和 Vision Transformer （ViTs） 等许多流行模型。
