[论文解读] Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

Lianghui Zhu, Bencheng Liao+4 more|arXiv (Cornell University)|Jan 17, 2024

Domain Adaptation and Few-Shot Learning被引用 384

一句话总结

Vision Mamba (Vim) 引入双向状态空间模型作为纯 SSM 视觉骨干，在高分辨率图像上以较低的计算和内存要求实现与 ViTs 相当的准确性。

ABSTRACT

Recently the state space models (SSMs) with efficient hardware-aware designs, i.e., the Mamba deep learning model, have shown great potential for long sequence modeling. Meanwhile building efficient and generic vision backbones purely upon SSMs is an appealing direction. However, representing visual data is challenging for SSMs due to the position-sensitivity of visual data and the requirement of global context for visual understanding. In this paper, we show that the reliance on self-attention for visual representation learning is not necessary and propose a new generic vision backbone with bidirectional Mamba blocks (Vim), which marks the image sequences with position embeddings and compresses the visual representation with bidirectional state space models. On ImageNet classification, COCO object detection, and ADE20k semantic segmentation tasks, Vim achieves higher performance compared to well-established vision transformers like DeiT, while also demonstrating significantly improved computation & memory efficiency. For example, Vim is 2.8$ imes$ faster than DeiT and saves 86.8% GPU memory when performing batch inference to extract features on images with a resolution of 1248$ imes$1248. The results demonstrate that Vim is capable of overcoming the computation & memory constraints on performing Transformer-style understanding for high-resolution images and it has great potential to be the next-generation backbone for vision foundation models. Code is available at https://github.com/hustvl/Vim.

研究动机与目标

• 为视觉任务提供一个纯状态空间模型骨干，以取代基于注意力的架构。
• 将双向状态空间建模和位置信息嵌入应用于视觉数据。
• 展示在高分辨率图像上计算与内存的高效性。
• 展示 ViM 在 ImageNet 分类及下游密集预测任务中的有效性。

提出的方法

• 采用基于 Mamba 的双向 SSM 块来处理图像补丁序列。
• 引入 Vim 块，应用带学习投影和门控的前向与后向 SSM。
• 为补丁令牌和分类令牌添加位置嵌入。
• 使用对 SRAM/HBM 内存友好的执行与重新计算以降低内存和 I/O。
• 提供一个具有 L 个 Vim 块、D 个隐藏维度和 E 个扩展维度的架构。
• 在 ImageNet、ADE20K 和 COCO 上将 Vim 与基于 ViT 的骨干和 SSM 基骨架进行比较。

实验结果

研究问题

• RQ1 纯 SSM 骨干是否能在常规基准上与基于 Transformer 的视觉模型相匹配或超越？
• RQ2 双向 SSM 建模是否为密集预测提供了足够的全局上下文和空间感知？
• RQ3 在高分辨率图像上，Vim 相对于 DeiT 的效率提升（速度与内存）有哪些？
• RQ4 设计选择（分类令牌策略、双向配置）如何影响分类和分割任务的性能？

主要发现

• Vim 的特征提取在 1248x1248 图像上比 DeiT 快 2.8 倍，GPU 内存减少 86.8%。
• 在 ImageNet 分类任务中，Vim 在多个模型尺度上均表现优于 DeiT。
• 带有向后通路的双向 SSM 及 Conv1d 增强比单向设置在分割与分类上表现更好。
• 在 COCO 上，Vim-Ti 在 AP、框和掩码的 AP 上均优于 DeiT-Ti，显示出更强的远程上下文学习能力。
• Vim 在没有 2D 先验的情况下实现高分辨率序列视觉表示学习，在若干设置中以更少的参数保持具有竞争力甚至更高的准确性。