[论文解读] VMamba: Visual State Space Model

Yue Liu, Yunjie Tian+7 more|arXiv (Cornell University)|Jan 18, 2024

Visual Attention and Saliency Detection被引用 358

一句话总结

VMamba 引入了一个具有线性复杂度和全局感受野的二维视觉状态空间模型，使用 Cross-Scan Module 实现多方向的二维信息整合，在 ImageNet、COCO 和 ADE20K 上达到有竞争力的性能。

ABSTRACT

Designing computationally efficient network architectures remains an ongoing necessity in computer vision. In this paper, we adapt Mamba, a state-space language model, into VMamba, a vision backbone with linear time complexity. At the core of VMamba is a stack of Visual State-Space (VSS) blocks with the 2D Selective Scan (SS2D) module. By traversing along four scanning routes, SS2D bridges the gap between the ordered nature of 1D selective scan and the non-sequential structure of 2D vision data, which facilitates the collection of contextual information from various sources and perspectives. Based on the VSS blocks, we develop a family of VMamba architectures and accelerate them through a succession of architectural and implementation enhancements. Extensive experiments demonstrate VMamba's promising performance across diverse visual perception tasks, highlighting its superior input scaling efficiency compared to existing benchmark models. Source code is available at https://github.com/MzeroMiko/VMamba.

研究动机与目标

• 提出一个将全局感受野与线性复杂度相结合的视觉基础模型。
• 利用带有选择性扫描方法的状态空间模型（SSMs）来在视觉任务中实现动态权重。
• 通过 Cross-Scan Module (CSM) 解决将 SSM 应用于 2D 图像时的方向敏感性。
• 将 VMamba 开发为可扩展的图像分类、目标检测与语义分割骨干网络。

提出的方法

• 在 Visual State Space (VSS) 块中采用 2D Selective Scan (SS2D)。
• 引入 Cross-Scan Module (CSM)，从四个角对 2D 特征图进行扫描，以在实现线性复杂度的同时确保全局感受野。
• 将连续状态空间动力学离散化以用于实际深度学习应用（矩阵指数和一阶近似）。
• 将 VSS 块堆叠成分层的 VMamba 骨干网络，采用与 ViT/CNN 骨干相似的阶段性下采样。
• 在 ImageNet-1K 上训练 VMamba 的变体（Tiny、Small、Base），并在 COCO 和 ADE20K 上进行检测与分割评估。

实验结果

研究问题

• RQ1 基于状态空间的视觉骨干网络是否能够在保持线性计算复杂度的同时实现全局感受野？
• RQ2 Cross-Scan Module 是否在不牺牲效率的前提下实现高效的 2D 信息整合？
• RQ3 在 ImageNet-1K、COCO 和 ADE20K 的不同尺度上，VMamba 与 CNNs 与 ViTs 的表现对比如何？
• RQ4 在输入分辨率放大/缩小时，VMamba 是否表现出鲁棒性并在训练后展示自适应的全局感受野？

主要发现

• VMamba 的各变体在 ImageNet-1K 的 top-1 准确率与主流基线相当，FLOPs 相近（例如 VMamba-T 在 4.5G FLOPs 时达到 82.2%）。
• 在 COCO 的 1x 和 3x MS 调度下，VMamba-T/S/B 在多种设置中均在框检测和掩码 AP 上超过 Swin 与 ConvNeXt 基线。
• 在 ADE20K 上使用 512x512 输入时，VMamba-S 与 VMamba-B 获得更高的 mIoU，VMamba-T 在 512x512 与多尺度测试下实现强分割性能。
• Cross-Scan Module 产生全局有效感受野（ERF），并实现十字形、长程依赖，训练后 ERF 变为全局。
• VMamba 随输入分辨率增加呈线性 FLOPs 增长，同时保持全局感受野，与需要二次复杂度的 ViTs 不同。