Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model
TL;DR 精炼摘要
Vision Mamba (Vim) 创新性地提出了一个基于双向状态空间模型(SSM)的高效通用视觉骨干网络,旨在克服Vision Transformer自注意力机制在高分辨率图像处理中面临的计算与内存瓶颈,并证明视觉表示学习无需依赖自注意力。该方法通过引入位置嵌入处理视觉数据的空间敏感性,并利用双向Mamba块高效压缩视觉信息。实验表明,Vim在ImageNet分类、目标检测及语义分割等任务上,不仅超越了传统ViT(如DeiT)的性能,更显著提升了计算速度(快2.8倍)和内存效率(节省86.8% GPU内存),展现了其成为下一代视觉基础模型骨干的巨大潜
摘要
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 faster than DeiT and saves 86.8% GPU memory when performing batch inference to extract features on images with a resolution of 12481248. 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.
思维导图
论文精读
中文精读
1. 论文基本信息 (Bibliographic Information)
- 标题 (Title): Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model (Vision Mamba: 基于双向状态空间模型的高效视觉表示学习)
- 作者 (Authors): Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, Xinggang Wang. 作者主要来自华中科技大学(HUST)等研究机构。
- 发表期刊/会议 (Journal/Conference): 这篇论文最初以预印本 (Preprint) 形式发布在 arXiv 上。arXiv 是一个广泛使用的学术论文预印本平台,允许研究人员在同行评审前分享他们的工作。
- 发表年份 (Publication Year): 2024年 (根据 arXiv 上的提交记录)
- 摘要 (Abstract): 近期,具备高效硬件感知设计的状态空间模型 (SSMs),特别是 Mamba 模型,在长序列建模方面展现了巨大潜力。因此,构建一个完全基于 SSMs 的高效通用视觉骨干网络成为一个引人注目的研究方向。然而,由于视觉数据的位置敏感性以及视觉理解对全局上下文的需求,使用 SSMs 表示视觉数据具有挑战性。本文证明了视觉表示学习并非必须依赖自注意力机制,并提出了一种新的通用视觉骨干网络——Vision Mamba (Vim)。Vim 通过位置嵌入来标记图像序列,并利用双向状态空间模型来压缩视觉表示。在 ImageNet 分类、COCO 目标检测和 ADE20k 语义分割等任务上,Vim 不仅取得了比成熟的 Vision Transformer (如 DeiT) 更高的性能,同时在计算和内存效率上也有显著提升。例如,在处理分辨率为 的图像进行批量推理时,Vim 的速度是 DeiT 的2.8倍,并节省了86.8%的 GPU 内存。这些结果表明,Vim 能够克服 Transformer 在处理高分辨率图像时的计算和内存瓶颈,有潜力成为下一代视觉基础模型的骨干网络。
- 原文链接 (Source Link):
https://github.com/hustvl/Vim(代码和模型),papers/68e3c225c83c981895f3bd33/paper.pdf(论文PDF)。该论文目前处于预印本状态。
2. 整体概括 (Executive Summary)
- 研究背景与动机 (Background & Motivation - Why):
- 核心问题:
Vision Transformer (ViT)已成为计算机视觉领域的标准架构,其核心优势在于通过self-attention机制捕获数据依赖的全局上下文。然而,self-attention的计算和内存复杂度与输入序列长度(即图像中的patch数量)成二次方关系 (),这使得 ViT 在处理高分辨率图像时变得非常昂贵,面临巨大的计算和内存瓶颈。 - 重要性与挑战: 高分辨率图像在医学影像、遥感、计算病理学等领域至关重要。现有方法通常通过引入局部窗口(如
Swin Transformer)或下采样来降低计算成本,但这牺牲了 ViT 原生的全局、序列化建模能力。 - 切入点/创新思路: 近期在自然语言处理领域大获成功的
Mamba模型,作为一种状态空间模型 (SSM),展现了对长序列建模的线性复杂度 (O(N)) 和卓越性能。这启发了研究者:是否可以用Mamba完全替代self-attention,构建一个同样强大但更高效的通用视觉骨干网络? 然而,直接将 Mamba 应用于视觉存在两大挑战:1) Mamba 是单向的,不符合图像二维空间信息需要全局、无方向性理解的本质;2) Mamba 缺乏对位置信息的感知,而这在视觉任务中至关重要。
- 核心问题:
- 核心贡献/主要发现 (Main Contribution/Findings - What):
- 提出了 Vision Mamba (Vim) 模型: 这是首批探索纯
SSM架构作为通用视觉骨干网络的工作之一。Vim 成功地将 Mamba 的高效性迁移到了视觉领域,证明了自注意力并非视觉表示学习的必要条件。 - 设计了双向 Mamba 模块 (Bidirectional Mamba Block): 为了解决 Mamba 单向性的问题,Vim 创新地设计了双向处理机制。它同时从前向和后向扫描图像
patch序列,从而捕获全局上下文信息,使其更适合处理空间数据。 - 集成了位置信息: 通过引入位置嵌入 (
position embeddings),Vim 解决了SSM缺乏位置感知的问题,使其能够胜任分类、检测、分割等多种视觉任务,特别是对位置敏感的稠密预测任务。 - 实现了卓越的效率和性能: 实验证明,Vim 在性能上优于或持平于同等规模的
DeiT模型,同时在处理高分辨率图像时,其速度和内存效率远超DeiT,展现出作为下一代视觉基础模型骨干的巨大潜力。
- 提出了 Vision Mamba (Vim) 模型: 这是首批探索纯
3. 预备知识与相关工作 (Prerequisite Knowledge & Related Work)
- 基础概念 (Foundational Concepts):
- Vision Transformer (ViT): 一种将
Transformer架构应用于计算机视觉的模型。它将输入图像分割成一系列不重叠的图像块 (patches),将每个patch线性投影成一个向量(token),并像处理单词一样处理这些token序列。其核心是self-attention机制,允许每个patch与所有其他patch进行交互,从而捕获全局依赖关系。但其计算复杂度为 ,其中 是patch的数量。 - 状态空间模型 (State Space Model - SSM): 源自经典控制理论,SSM 通过一个隐状态 来描述系统如何从输入 演化到输出 。在深度学习中,现代 SSM(如 S4, Mamba)被离散化并设计成可以高效并行计算的神经网络层。它们擅长建模长程依赖,并且计算复杂度通常与序列长度成线性或近线性关系 (
O(N)或 )。 - Mamba: 一种创新的
SSM架构。其核心特点是引入了选择性机制,使得模型的参数(特别是状态转移矩阵)是输入依赖的(input-dependent),这使其能够像self-attention一样根据输入内容动态地聚焦于不同信息。此外,Mamba 采用了一种硬件感知的并行扫描算法,极大地提升了训练和推理效率。
- Vision Transformer (ViT): 一种将
- 前人工作 (Previous Works):
- 卷积神经网络 (ConvNets): 如
ResNet,是视觉领域的经典架构,通过卷积核提取局部特征,具有强大的平移不变性等归纳偏置。 - 视觉 Transformer (ViTs): 如
ViT、DeiT、Swin Transformer。DeiT通过知识蒸馏等技术提升了 ViT 的数据效率。Swin Transformer等分层架构通过在局部窗口内计算自注意力来降低计算量,并引入了层级结构,但牺牲了 ViT 纯粹的序列建模范式。 - 用于视觉的 SSM: 先前的工作要么将
SSM与CNN或Transformer结合构成混合架构 (e.g.,U-Mamba),要么将其应用于特定的视觉任务(如视频理解),但很少有工作探索构建一个纯粹基于 SSM 的通用视觉骨干网络。
- 卷积神经网络 (ConvNets): 如
- 技术演进 (Technological Evolution): 视觉骨干网络经历了从
ConvNets(强调局部归纳偏置)到ViTs(强调全局、数据依赖的上下文建模)的演进。Vim 的出现标志着一个新的方向:探索SSM作为第三种主流范式,它试图在保持ViT全局建模能力的同时,实现远超ViT的计算效率。 - 差异化分析 (Differentiation):
- 与 ViT/DeiT 相比: Vim 的核心区别在于用双向选择性 SSM 替换了
self-attention模块。这使得 Vim 在保持全局上下文建模能力的同时,将计算复杂度从二次方降至线性。 - 与其他视觉 SSM 模型相比: Vim 是一个纯 SSM 架构,旨在成为一个像
ResNet或ViT一样的通用骨干网络,而不是为特定任务设计的混合模型。 - 与原始 Mamba 相比: Vim 的核心创新在于双向建模和位置嵌入的引入,使其能够处理非因果、位置敏感的图像数据,而原始 Mamba 主要为单向的文本序列设计。
- 与 ViT/DeiT 相比: Vim 的核心区别在于用双向选择性 SSM 替换了
4. 方法论 (Methodology - Core Technology & Implementation Details)
Vim 的核心目标是构建一个高效且强大的视觉骨干网络,其整体架构和核心模块设计如下。
-
方法原理 (Methodology Principles):
- 核心思想是利用
Mamba的选择性SSM机制来模拟self-attention的数据依赖性,同时通过双向扫描来捕获图像的全局空间信息,并借助位置嵌入来提供必要的空间位置感知。
- 核心思想是利用
-
方法步骤与流程 (Steps & Procedures): 参考下图,Vim 的整体流程与 ViT 非常相似,主要区别在于其编码器内部的设计。
该图像为示意图,展示了Vision Mamba(Vim)模型的整体架构及其编码器内部结构。左侧展示了输入图像经过分块、位置编码、线性投影形成带有位置标记的patch tokens,随后输入Vision Mamba编码器进行处理;右侧为编码器细节,包含归一化层及双向卷积和双向状态空间模型(SSM)模块,前向和后向信息流相互融合,最终生成输出特征。- 图像分块与线性投影 (Patching & Projection): 将输入图像 分割成一系列不重叠的 的图像块 (
patches)。将这些二维patches展平并线性投影成一维的token序列。 - 添加特殊 Token 和位置嵌入:
- 在序列的开头(或中间,根据消融实验结果)插入一个可学习的
class token,用于最终的图像分类任务。 - 为整个
token序列添加可学习的位置嵌入 (position embeddings),以使模型能够感知每个patch的空间位置。
- 在序列的开头(或中间,根据消融实验结果)插入一个可学习的
- Vision Mamba 编码器 (Vim Encoder):
- 处理后的
token序列被送入一个由 L 个Vim Block堆叠而成的编码器中。 - 每个
Vim Block内部都包含一个双向 Mamba 模块,这是 Vim 的核心。
- 处理后的
- Vim Block 详解 (右图):
- 输入归一化:
token序列首先经过一个归一化层 (Norm)。 - 线性投影: 归一化后的
token被线性投影到两个分支:主干分支 x 和门控分支 z。 - 双向处理: x 分支被并行地送入两个子模块进行处理:
- 前向路径 (Forward Path): 序列 x 依次通过一个一维卷积 (
Forward Conv1d) 和一个前向的选择性状态空间模型 (Forward SSM)。 - 后向路径 (Backward Path): 序列 x 被逆序送入一个独立的一维卷积 (
Backward Conv1d) 和一个后向的选择性状态空间模型 (Backward SSM)。
- 前向路径 (Forward Path): 序列 x 依次通过一个一维卷积 (
- 信息融合与门控:
- 前向和后向路径的输出
y_forward和y_backward分别与门控分支 z(经过SiLU激活函数)相乘。 - 两个被门控后的输出相加,从而融合了双向的上下文信息。
- 前向和后向路径的输出
- 输出与残差连接: 融合后的结果经过一个最终的线性投影层,然后通过残差连接与该
Vim Block的输入相加,形成最终输出。
- 输入归一化:
- 图像分块与线性投影 (Patching & Projection): 将输入图像 分割成一系列不重叠的 的图像块 (
-
数学公式与关键细节 (Mathematical Formulas & Key Details):
- SSM 离散化:
SSM 模型首先定义一个连续系统,然后通过一个时间尺度参数 将其离散化。离散化后的状态更新和输出计算如下:
其中:
- 是在时间步 的隐状态。
- 是在时间步 的输入。
- 是在时间步 的输出。
- 和 是离散化的状态矩阵,它们由连续矩阵 , 和时间尺度参数 计算得出。 是投影矩阵。
- 计算复杂度:
对于一个长度为 、维度为 的序列,
self-attention和SSM的计算复杂度分别为: 其中 是 SSM 的状态维度(一个很小的常数,默认为16)。- 分析:
self-attention的复杂度与序列长度 成二次方关系 (),而SSM的复杂度与 成线性关系 (O(M D N))。这正是 Vim 在处理高分辨率图像(即 很大)时高效的关键原因。
- 分析:
- SSM 离散化:
SSM 模型首先定义一个连续系统,然后通过一个时间尺度参数 将其离散化。离散化后的状态更新和输出计算如下:
其中:
5. 实验设置 (Experimental Setup)
- 数据集 (Datasets):
- ImageNet-1K: 用于图像分类任务。包含约128万张训练图像和5万张验证图像,共1000个类别。是评估视觉骨干网络性能的标准基准。
- ADE20K: 用于语义分割任务。是一个场景解析数据集,包含2万多张图像和150个语义类别。
- COCO 2017: 用于目标检测和实例分割任务。包含11.8万张训练图像和5千张验证图像,涵盖80个常见物体类别。
- 评估指标 (Evaluation Metrics):
- Top-1 Accuracy: 用于 ImageNet 分类,表示模型预测的最可能类别与真实标签一致的百分比。
- mIoU (mean Intersection over Union): 用于语义分割,计算所有类别的交并比(预测区域与真实区域的交集/并集)的平均值,是衡量分割质量的核心指标。
- AP (Average Precision): 用于目标检测和实例分割,综合衡量模型的查准率和查全率。论文中报告了多个 AP 指标,如 (边界框精度) 和 (掩码精度)。
- 对比基线 (Baselines):
- DeiT (Data-efficient image Transformers): 主要的对比基线。DeiT 是一个经过高度优化的 ViT 变体,在性能和训练效率上都很有代表性。
- ConvNets: 如
ResNet系列,作为经典的卷积网络基线。 - 其他 ViT 变体和 SSM 模型: 如
S4ND-ViT-B。
6. 实验结果与分析 (Results & Analysis)
-
核心结果分析 (Core Results Analysis):
该图像为三部分图表,分别比较了ViM-Ti与DeiT-Ti模型在多个任务上的性能。左图(a)显示ViM-Ti在分类、语义分割、目标检测和实例分割上的准确率均优于DeiT-Ti。中图(b)展示不同分辨率下两模型推理速度,ViM-Ti速度明显更快,最高达2.8倍提升。右图(c)对比GPU内存占用,ViM-Ti显著节省内存,分辨率1248时DeiT-Ti因内存溢出而无法运行。整体说明ViM模型在准确率、速度和内存效率上均表现更优。-
分类、检测、分割性能全面超越 DeiT (上图 a):
- 在 ImageNet 分类任务上,Vim-Ti, Vim-S, Vim-B 等不同尺寸的模型均取得了比对应 DeiT 模型更高的 Top-1 准确率 (见论文 Table 1)。
- 在 ADE20K 语义分割和 COCO 检测/分割任务上,Vim 同样全面优于 DeiT。这证明了 Vim 不仅在分类上强大,其学到的视觉表示也能很好地迁移到下游的稠密预测任务。
-
高分辨率下的效率优势巨大 (上图 b, c 和下图):
-
速度 (Speed): 随着输入图像分辨率的增加,Vim 的速度优势愈发明显。在 分辨率下,Vim 的推理速度是 DeiT 的 2.8倍。
-
内存 (Memory): Vim 的内存占用增长远比 DeiT 缓慢。在 分辨率下,DeiT 已经因为内存溢出 (
OOM) 而无法运行,而 Vim 仅需少量内存,节省了 86.8% 的 GPU 显存。 -
结论: 这种线性的计算和内存扩展性是 Vim 相对于 Transformer 的核心优势,使其成为处理高分辨率视觉任务的理想选择。
该图像为折线图,比较了DeiT和Vim模型在不同分辨率下的帧率(FPS,带对数刻度)。横轴为图像分辨率(512到1248),纵轴为FPS值。图中显示,Vim在所有分辨率下的FPS均高于DeiT,尤其在1248分辨率时,Vim的FPS约为1.70,而DeiT仅为1.25,表明Vim在该分辨率下实现了约2.8倍的速度提升。图左侧箭头标注“Faster”进一步强调了Vim的性能优势。
该图像为图表,展示了不同分辨率下DeiT与Vim两种模型的GPU显存占用情况。横轴为图像分辨率,纵轴为GPU显存(GB)。随着分辨率提升,DeiT显存需求显著上升,1248分辨率时达到OOM(显存溢出),而Vim显存增长较缓,1248分辨率仅需22.59GB,节省约73.2%显存,体现了Vim在高分辨率下更优的显存效率。
-
-
定性分析 (Qualitative Analysis): 下图展示了一个实例分割的例子。Vim-Ti 能够更准确、完整地识别出飞机的主体,而 DeiT-Ti 则出现了漏检。这直观地表明 Vim 学习到了更强大的视觉表示。
该图像为插图,展示了在同一场景上一架飞机的三张图片。左图为真实标签(GT),中图为DeiT-Ti模型的检测结果,右图为Vim-Ti模型的检测结果。Vim-Ti在飞机区域覆盖更完整且置信度更高(86%),显示其在视觉表示学习上的优势。
-
-
消融实验/参数分析 (Ablation Studies / Parameter Analysis):
- 双向设计的有效性 (Ablation on Bidirectional Design):
- 论文在 Table 4 中对比了多种双向策略。结果显示,简单的单向 Mamba (
None) 虽然在分类上表现尚可,但在对空间信息要求更高的语义分割任务上表现很差 (32.3 mIoU)。 - 最终采用的
Bidirectional SSM + Conv1d策略在分类 (73.9 Top-1 Acc) 和分割 (35.9 mIoU) 上均取得了最佳性能。这充分证明了双向建模对于视觉任务至关重要。
- 论文在 Table 4 中对比了多种双向策略。结果显示,简单的单向 Mamba (
- 分类设计的选择 (Ablation on Classification Design):
- 论文在 Table 5 中探索了不同
class token放置策略。结果发现,将class token放在序列中间 (Middle class token) 取得了最高的分类准确率 (76.1 Top-1 Acc)。 - 分析: 这可能是因为 SSM 的循环特性使得中间位置的
token能够最充分地聚合来自前向和后向两个方向的完整序列信息。
- 论文在 Table 5 中探索了不同
- 双向设计的有效性 (Ablation on Bidirectional Design):
7. 总结与思考 (Conclusion & Personal Thoughts)
-
结论总结 (Conclusion Summary):
- 本文成功提出了 Vision Mamba (Vim),一个基于纯状态空间模型的通用视觉骨干网络。
- 通过创新的双向状态空间建模和位置嵌入,Vim 克服了原始 Mamba 在应用于视觉任务时的局限性。
- 实验结果表明,Vim 在多个基准测试中性能优于强大的 DeiT 基线,同时在处理高分辨率图像时展现出压倒性的计算和内存效率优势。
- 这项工作有力地证明了
self-attention并非不可替代,高效的SSM架构有巨大潜力成为下一代视觉基础模型的基石。
-
局限性与未来工作 (Limitations & Future Work):
- 未来工作: 作者指出,Vim 的架构非常适合进行大规模无监督预训练,如掩码图像建模 (
Masked Image Modeling) 和多模态预训练 (如CLIP风格)。 - 应用前景: 基于 Vim 的高效率,可以很直接地将其应用于需要处理超高分辨率图像的下游任务,例如医学影像分析、遥感图像解译和长视频理解。
- 未来工作: 作者指出,Vim 的架构非常适合进行大规模无监督预训练,如掩码图像建模 (
-
个人启发与批判 (Personal Insights & Critique):
- 启发: 这篇论文最大的启发在于,它打破了
Transformer在高性能视觉模型中的“垄断地位”。它展示了一条全新的、以效率为导向的技术路线,即利用SSM来实现与self-attention相似的数据依赖性全局建模,但成本要低得多。这对于推动视觉模型在资源受限设备或超大规模数据上的应用具有重要意义。 - 批判性思考:
- 归纳偏置的权衡: Vim 像 ViT 一样,保留了较少的图像特有归纳偏置(如平移不变性),这可能使其在数据量较少时仍然需要大量数据进行预训练。虽然性能优于 DeiT,但与精心设计的 ConvNets 或混合架构(如
Swin Transformer)相比,其在中小规模数据集上的数据效率可能仍是一个值得深入研究的问题。 - 双向扫描的本质: 尽管双向扫描能捕获全局信息,但它本质上仍是对一维序列的操作。与
2D窗口或卷积相比,它在捕捉局部二维空间结构方面的能力如何,仍需更深入的分析。论文的附录通过构建分层版本的Hier-Vim并取得了优于Swin-Transformer的结果,在一定程度上回应了这个问题,证明了其潜力。 - 并发工作: 论文提到了一个并发工作
VMamba。VMamba采用了多方向扫描(沿行、列、之字形扫描)来更好地捕捉二维信息,而 Vim 更侧重于保持纯粹的序列建模范式。这两种不同思路的优劣对比将是未来一个有趣的研究点。
- 归纳偏置的权衡: Vim 像 ViT 一样,保留了较少的图像特有归纳偏置(如平移不变性),这可能使其在数据量较少时仍然需要大量数据进行预训练。虽然性能优于 DeiT,但与精心设计的 ConvNets 或混合架构(如
- 启发: 这篇论文最大的启发在于,它打破了
相似论文推荐
基于向量语义检索推荐的相关论文。