Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

📄 arXiv: 2401.09417v3 📥 PDF

作者: Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, Xinggang Wang

分类: cs.CV, cs.LG

发布日期: 2024-01-17 (更新: 2024-11-14)

备注: Vision Mamba (Vim) is accepted by ICML 2024. Code is available at https://github.com/hustvl/Vim

🔗 代码/项目: GITHUB

💡 一句话要点

提出Vim以解决视觉表示学习中的计算与内存效率问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 状态空间模型 视觉表示学习 双向Mamba 计算效率 内存优化 图像处理 深度学习

📋 核心要点

  1. 现有方法在视觉表示学习中依赖自注意力机制，导致计算和内存效率低下，尤其在处理高分辨率图像时。
  2. 本文提出的Vim网络通过双向Mamba模块，结合位置嵌入和状态空间模型，减少对自注意力的依赖，从而提高效率。
  3. 实验结果表明，Vim在多个视觉任务上超越了DeiT，且在批量推理时速度提升2.8倍，GPU内存节省86.8%。

📝 摘要（中文）

近年来，具有高效硬件感知设计的状态空间模型（SSMs），如Mamba深度学习模型，在长序列建模中展现出巨大潜力。然而，基于SSMs构建高效通用的视觉骨干网络面临挑战，主要由于视觉数据的位置信息敏感性和对全局上下文的理解需求。本文提出了一种新的通用视觉骨干网络Vim，利用双向Mamba模块，通过位置嵌入标记图像序列，并压缩视觉表示。Vim在ImageNet分类、COCO目标检测和ADE20k语义分割任务中表现优异，相较于成熟的视觉变换器如DeiT，Vim在计算和内存效率上显著提升，具有成为下一代视觉基础模型骨干的潜力。

🔬 方法详解

问题定义：本文旨在解决现有视觉表示学习方法在计算和内存效率上的不足，尤其是在处理高分辨率图像时，传统的自注意力机制导致了性能瓶颈。

核心思路：论文提出了一种新的视觉骨干网络Vim，利用双向Mamba模块，结合位置嵌入来标记图像序列，从而压缩视觉表示，减少对自注意力的依赖。

技术框架：Vim的整体架构包括双向Mamba模块和位置嵌入模块，前者负责视觉信息的处理与压缩，后者则确保位置信息的有效传递。

关键创新：Vim的核心创新在于引入双向状态空间模型，显著提高了计算和内存效率，与传统的视觉变换器方法相比，Vim在处理高分辨率图像时表现出更优的性能。

关键设计：在网络结构设计上，Vim采用了位置嵌入来增强位置信息的表达，同时在损失函数和参数设置上进行了优化，以确保模型在不同任务中的适应性和效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Vim在多个视觉任务中表现优异，尤其在ImageNet分类、COCO目标检测和ADE20k语义分割任务上，性能超越了DeiT，且在批量推理时速度提升2.8倍，GPU内存节省86.8%，展示了其在高效视觉表示学习中的巨大潜力。

🎯 应用场景

该研究的潜在应用领域包括计算机视觉、自动驾驶、智能监控等，尤其在需要处理高分辨率图像的场景中，Vim的高效性将显著提升系统的响应速度和处理能力。未来，Vim有望成为视觉基础模型的主流选择，推动相关技术的发展。

📄 摘要（原文）

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$\times$ faster than DeiT and saves 86.8% GPU memory when performing batch inference to extract features on images with a resolution of 1248$\times$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.