Vision Mamba: 基于双向状态空间模型的高效视觉表征学习

基本信息

• 标题：Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model
• 作者：Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, Xinggang Wang
• arXiv 编号：arXiv:2401.09417（v2）
• 提交日期：2024 年 1 月 17 日（v1）；本版本修订于 2024 年 2 月 10 日（v2）；最新版本 v3 于 2024 年 11 月 14 日
• 主题分类：Computer Vision and Pattern Recognition (cs.CV)；Machine Learning (cs.LG)
• 备注：Work in progress，代码见 https://github.com/hustvl/Vim

摘要

近年来，具有高效硬件感知设计的状态空间模型（State Space Models, SSMs），即 Mamba 深度学习模型，在长序列建模方面展现出巨大潜力。与此同时，完全基于 SSM 构建高效且通用的视觉骨干网络是一个颇具吸引力的方向。然而，由于视觉数据对位置敏感，且视觉理解需要全局上下文信息，用 SSM 表征视觉数据颇具挑战。

本文表明，视觉表征学习并不一定依赖自注意力机制，并提出了一种全新的通用视觉骨干网络——Vim（Vision Mamba）。它采用双向 Mamba 模块，为图像序列添加位置嵌入，并通过双向状态空间模型压缩视觉表征。

在 ImageNet 分类、COCO 目标检测和 ADE20k 语义分割任务上，Vim 相比 DeiT 等成熟的视觉 Transformer 取得了更高的性能，同时显著提升了计算与内存效率。例如，在对 1248×1248 分辨率的图像进行批量推理以提取特征时，Vim 比 DeiT 快 2.8 倍，并节省了 86.8% 的 GPU 内存。结果表明，Vim 能够突破 Transformer 式高分辨率图像理解所面临的计算与内存限制，有潜力成为下一代视觉基础模型骨干网络。代码已开源于 https://github.com/hustvl/Vim。

主要贡献

• 提出通用视觉骨干网络 Vim，完全基于双向状态空间模型（Mamba），不依赖自注意力机制。
• 通过位置嵌入标记图像序列，利用双向 SSM 压缩视觉表征。
• 在 ImageNet 分类、COCO 目标检测、ADE20k 语义分割任务上优于 DeiT 等视觉 Transformer。
• 在高分辨率图像特征提取上表现出显著的效率优势（比 DeiT 快 2.8 倍，节省 86.8% GPU 内存）。

相关链接

• 代码：https://github.com/hustvl/Vim
• DOI：https://doi.org/10.48550/arXiv.2401.09417