Vision Mamba论文精读笔记

这篇博客主要针对Vision Mamba 论文进行精读，包含全文翻译以及部分内容注解。

读者最好有SSM以及Mamba的前期基础，便于理解。

论文链接：[2401.09417] Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

开源代码：GitHub - hustvl/Vim: [ICML 2024] Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

摘要

最近，具有高效硬件意识设计的状态空间模型（SSMs），即 Mamba 深度学习模型，在长序列建模方面展现出巨大的潜力。同时，基于 SSMs 构建高效且通用的视觉骨干网络成为一个吸引人的方向。然而，由于视觉数据的位置信息敏感性和理解视觉内容所需的全局上下文，使用 SSMs 表示视觉数据是一项挑战。

在本文中，我们表明，视觉表示学习对自注意力的依赖并不是必要的，并提出了一种新的通用视觉骨干网络——双向 Mamba 块 (bidirectional Mamba blocks)——Vim。该网络通过位置嵌入标记图像序列，并利用双向状态空间模型压缩视觉表示。

在 ImageNet 分类、COCO 目标检测和 ADE20k 语义分割任务中，Vim 相较于知名的视觉变换器（如 DeiT）取得了更高的性能，同时在计算和内存效率上也显著提升。例如，在对分辨率为 1248×1248 的图像进行批量推理以提取特征时，Vim 的速度比 DeiT 快 2.8 倍，并节省了 86.8% 的 GPU 内存。结果表明，Vim 能够克服 Transformer 风格的高分辨率图像理解在计算和内存上的限制，具有成为下一代视觉基础模型骨干网络的巨大潜力。

一、引言

SSM的背景和发展

最近的研究进展引发了对状态空间模型（SSM）的大量关注。状态空间模型起源于经典的卡尔曼滤波模型（Kalman, 1960），现代 SSM 在捕捉长程依赖性方面表现出色，并得益于并行训练。一些基于 SSM 的方法，如线性状态空间层（LSSL）（Gu et al., 2021b）、结构化状态空间序列模型（S4）（Gu et al., 2021a）、对角状态空间（DSS）（Gupta et al., 2022）和 S4D（Gu et al., 2022），被提出用于处理各种任务和模态的序列数据，特别是在建模长程依赖性方面。

2D SSM和Mamba的进展

由于卷积计算和近线性计算，它们在处理长序列时效率较高。2-D SSM（Baron et al., 2023）、SGConvNeXt（Li et al., 2022b）和 ConvSSM（Smith et al., 2023a）将 SSM 与 CNN 或 Transformer 架构结合起来，以处理 2-D 数据。

最近的工作 Mamba（Gu & Dao, 2023）将时间变化参数纳入 SSM，并提出了一种硬件感知算法，以实现非常高效的训练和推理。Mamba 的优越扩展性能表明，它是语言建模中 Transformer 的一种有前景的替代方案。然而，目前尚未探索出用于处理视觉数据（如图像和视频）的通用纯 SSM 基础骨干网络。

Vision Transformer 的成功与局限

视觉变换器（ViTs）在视觉表示学习中取得了巨大成功，尤其是在大规模自监督预训练和下游任务中表现优异。

与卷积神经网络相比，ViT 的核心优势在于可以通过自注意力为每个图像块提供数据/块依赖的全局上下文。这与卷积网络不同，后者对所有位置使用相同的参数，即卷积滤波器。

另一个优势是通过将图像视为一系列块而进行模态无关建模，从而没有 2D 归纳偏置（Inductive Bias），这使其成为多模态应用的首选架构。

什么是归纳偏置(Inductive Bias)？

归纳偏置会促使学习算法优先考虑具有某些属性的解。

在CNN中，常见归纳偏置为局部相关性（Local Connectivity）和平移不变性（Translation Invariance）。

- 局部相关性：图像中的信息往往是局部相关的，即一个像素的值通常与其周围的像素密切相关。CNN通过卷积核（filter）在局部区域内提取特征，能减少参数量，高效地捕获局部特征。

- 平移不变性：图像中的特征（如物体）在图像中的位置不影响其类别或属性。CNN 中的卷积操作和池化操作天然支持平移不变性。例如，卷积核在图像上滑动（共享权重），能够检测特定特征的位置无关性。提高了模型对位置信息的鲁棒性，使其在不同位置检测相同对象。

为什么没有归纳偏置是一种优势？

- 任务依赖性：虽然CNN的归纳偏置（如局部相关性、平移不变性）使其在处理图像数据时表现出色，但对于非图像数据（如序列数据）时表现不佳；

- 灵活性不足：归纳偏置限制了模型的表达能力，使其无法捕捉更复杂或全局的依赖关系。CNN的局部卷积操作可能无法有效捕捉图像中的长距离依赖关系。

ViT通过将图像视为一系列块（patches）并进行模态无关建模，减少了对2D归纳偏置的依赖。这种设计带来了以下优势：

全局上下文建模：ViT通过自注意力机制（self-attention）为每个图像块提供数据依赖的全局上下文，而不是像CNN那样使用固定的卷积核。
这使得ViT能够更好地捕捉图像中的长距离依赖关系。

模态无关性：ViT将图像视为序列数据，这种设计使其能够更灵活地处理多模态任务（如文本-图像联合建模）。
例如，ViT可以直接处理文本、图像甚至其他类型的数据，而无需针对每种模态设计特定的模型结构。

可扩展性：ViT在大规模数据集上表现出色，尤其是在自监督预训练任务中。
由于减少了归纳偏置，ViT能够从大量数据中学习更通用的特征表示。

同时，Transformer 中的自注意力机制在处理长程视觉依赖性（例如处理高分辨率图像）时在速度和内存使用方面带来了挑战。

Vision Mamba (Vim) 的提出

受到 Mamba 在语言建模中成功的启发，我们认为可以将这种成功从语言转移到视觉，即设计一种具有先进 SSM 方法的通用高效视觉骨干网络。然而，Mamba 面临两个挑战，即单向建模和缺乏位置意识。

为了解决这些挑战，我们提出了视觉 Mamba（Vim）模型，该模型结合了双向 SSM 用于数据依赖的全局视觉上下文建模，并使用位置嵌入实现位置感知的视觉识别。

我们首先将输入图像拆分为多个块，并将它们线性投影为向量传递给 Vim。
图像块被视为 Vim 块中的序列数据，通过提出的双向选择状态空间有效压缩视觉表示。
此外，Vim 块中的位置嵌入提供了空间信息的感知，使 Vim 在密集预测任务中更具鲁棒性。

在当前阶段，我们在使用 ImageNet 数据集的有监督图像分类任务上训练 Vim 模型，然后使用预训练的 Vim 作为骨干进行下游密集预测任务的序列视觉表示学习，即语义分割、目标检测和实例分割。与 Transformer 类似，Vim 可以在大规模无监督视觉数据上进行预训练，以获得更好的视觉表示。得益于 Mamba 更好的效率，Vim 的大规模预训练可以以更低的计算成本实现。

Vim的贡献与性能

与其他用于视觉任务的 SSM 基础模型相比，Vim 是一种纯 SSM 基础的方法，并以序列方式对图像建模，这使其更有前景作为通用高效的骨干网络。由于其具备位置感知的双向压缩建模，Vim 是首个能够处理密集预测任务的纯 SSM 基础模型。与最具说服力的基于 Transformer 的模型（即 DeiT）（Touvron et al., 2021a）相比，Vim 在 ImageNet 分类上取得了更优越的性能。此外，Vim 在高分辨率图像的 GPU 内存和推理时间方面更为高效。这种在内存和速度方面的高效性使 Vim 能够直接进行序列视觉表示学习，而无需依赖于 2D 先验（如 ViTDet（Li et al., 2022c）中的 2D 局部窗口）来处理高分辨率视觉理解任务，同时在准确性上超越了 DeiT。

我们的主要贡献可以总结如下：

我们提出了视觉 Mamba（Vim），该模型结合了双向 SSM 用于数据依赖的全局视觉上下文建模和位置嵌入用于位置感知的视觉理解。
Vim 在不需要注意力机制的情况下，具备与 ViT 相同的建模能力，同时仅具有亚二次时间计算和线性内存复杂度。具体而言，Vim 在对分辨率为 1248×1248 的图像执行批量推理以提取特征时，比 DeiT 快 2.8 倍，并节省了 86.8% 的 GPU 内存。
我们在 ImageNet 分类和密集预测下游任务上进行了广泛实验。结果表明，Vim 相较于成熟且高度优化的纯视觉 Transformer（即 DeiT）实现了更优越的性能。

二、相关工作

通用视觉骨干网络的架构

在早期，卷积神经网络是计算机视觉领域的事实标准网络设计。许多卷积神经网络架构被提出作为各种视觉应用的骨干网络。开创性的工作——视觉Transformer改变了这一格局。它将图像视为一系列展平的2D图像块，并直接应用纯Transformer架构。视觉Transformer在图像分类任务中的惊人表现及其扩展能力激励了许多后续工作。一部分工作专注于通过将2D卷积先验引入视觉Transformer来设计混合架构。PVT提出了一种金字塔结构的Transformer。Swin Transformer在滑动窗口内应用自注意力机制。另一部分工作则专注于通过更先进的设置改进传统的2D卷积网络。ConvNeXt回顾了设计空间，并提出了一种纯卷积网络，能够像视觉Transformer及其变体一样具有扩展性。RepLKNet提出通过扩大现有卷积网络的核尺寸来带来改进。

尽管这些主导的后续工作在ImageNet和各种下游任务中通过引入2D先验展示了卓越的性能和更高的效率，但随着大规模视觉预训练和多模态应用的兴起，vanilla Transformer风格的模型重新回到了计算机视觉的中心舞台。其更大的建模能力、统一的多模态表示以及对自监督学习的友好性等优势使其成为首选架构。然而，由于Transformer的二次复杂度，视觉token的数量受到限制。有许多工作致力于解决这一长期存在的突出挑战，但其中很少专注于视觉应用。最近，LongViT通过扩张注意力机制构建了一种高效的Transformer架构，用于计算病理学应用。LongViT的线性计算复杂度使其能够编码极长的视觉序列。

在本工作中，我们从Mamba中汲取灵感，探索构建一种纯基于状态空间模型的通用视觉骨干网络，不使用注意力机制，同时保留视觉Transformer的序列化、模态无关的建模优势。

长序列建模的状态空间模型

（Gu et al., 2021a）提出了一种结构化状态空间序列（S4）模型，作为CNN或Transformer的替代方案，用于建模长程依赖关系。其在序列长度上线性扩展的潜力吸引了进一步的探索。（Wang et al., 2022）提出了双向门控SSM，以在不使用注意力机制的情况下复现BERT（Devlin et al., 2018）的结果。（Smith et al., 2023b）提出了一种新的S5层，通过将MIMO SSM和高效并行扫描引入S4层。（Fu et al., 2023）设计了一种新的SSM层H3，几乎填补了SSM与Transformer注意力在语言建模中的性能差距。（Mehta et al., 2023）在S4的基础上构建了门控状态空间层，通过引入更多门控单元来提高表达能力。最近，（Gu & Dao, 2023）提出了一种数据依赖的SSM层，并构建了通用语言模型骨干网络Mamba，该模型在大规模真实数据上优于Transformer，并在序列长度上具有线性扩展性。

在本工作中，我们探索将Mamba的成功迁移到视觉领域，即构建一种纯基于SSM的通用视觉骨干网络，不使用注意力机制。

视觉应用中的状态空间模型

（Islam & Bertasius, 2022）使用1D S4来处理视频分类中的长程时间依赖关系。（Nguyen et al., 2022）进一步将1D S4扩展到处理包括2D图像和3D视频在内的多维数据。（Islam et al., 2023）结合了S4和自注意力的优势，构建了TranS4mer模型，在电影场景检测中取得了最先进的性能。（Wang et al., 2023a）在S4中引入了一种新颖的选择性机制，大大提高了S4在长视频理解中的性能，同时显著降低了内存占用。（Yan et al., 2023）用更可扩展的基于SSM的骨干网络取代了注意力机制，以生成高分辨率图像并在可承受的计算量下处理细粒度表示。（Ma et al., 2024）提出了U-Mamba，一种混合的CNN-SSM架构，用于处理生物医学图像分割中的长程依赖关系。

上述工作（Xing et al., 2024; Ma et al., 2024; Yan et al., 2023; Wang et al., 2023a; Islam et al., 2023; Nguyen et al., 2022; Islam & Bertasius, 2022）要么将SSM应用于特定的视觉应用，要么通过将SSM与卷积或注意力机制结合来构建混合架构。

与它们不同，我们构建了一种纯基于SSM的模型，可以作为通用视觉骨干网络。值得注意的是，VMamba（Liu et al., 2024）是与我们方法同时期的工作，通过将Mamba与多方向扫描和分层网络架构结合，在视觉识别中展示了令人印象深刻的结果。相比之下，Vim主要专注于视觉序列学习，并拥有统一的多模态数据表示。

VMamba论文精读链接：VMamba论文精读笔记_vmamba模型-CSDN博客

三、方法

Vision Mamba (Vim) 的目标是将先进的状态空间模型（State Space Model, SSM），即 Mamba（Gu & Dao, 2023），引入计算机视觉领域。本节首先介绍 SSM 的基础知识，随后概述 Vim 的设计，并详细描述 Vim 模块如何处理输入 token 序列，最后分析 Vim 的效率。

3.1 基础知识

这部分内容可看：

从SSM到S4_state space model-CSDN博客

Mamba理论入门_mamba一次性完成前向传播-CSDN博客

基于 SSM 的模型，如结构化状态空间序列模型（S4）和 Mamba，灵感来源于连续系统。该系统通过隐藏状态 $h(t)\in \mathbb{R}^N$ 将一维函数或序列 $x(t)\in \mathbb{R}\rightarrow y(t)\in \mathbb{R}$ 进行映射。

系统使用 $A\in \mathbb{R}^{N\times N}$ 作为演化参数， $B\in \mathbb{R}^{N\times 1}$ 和 $C\in \mathbb{R}^{1\times N}$ 作为投影参数。连续系统的公式如下：

$h{}'(t)=Ah(t)+Bx(t),y(t)=Ch(t)$

S4 和 Mamba 是连续系统的离散版本，引入时间尺度参数 $\triangle$ 将连续参数 $A,B$ 转换为离散参数 $\bar{A},\bar{B}$ 。常用的转换方法是零阶保持（ZOH），定义如下：

零级保持

零阶保持(Zero-Order Hold,ZOH)用于将离散时间信号转换为连续时间信号。这种方法假设输入信号在每个采样周期内保持恒定不变，即在每个采样点之后，直到下一个采样点到来之前，信号的值不再变化。

常用离散化方法

主要目的是将连续时间系统（通常由微分方程描述）转换为离散时间系统（由差分方程描述），以便在计算机或数字系统中进行数值计算或实现。

欧拉法（Euler Method）

欧拉法是一种简单的数值积分方法，用于将连续时间微分方程离散化。
它通过将导数近似为有限差分，将连续时间系统转换为离散时间系统。
例如，对于连续时间微分方程 $\frac{dx(t)}{dt} = f(x(t),t)$ ，欧拉法的离散形式为： $x_{k+1}=k_k + \triangle t\cdot f(x_k,t_k)$
其中，Δt是时间步长， $x_k$ 是离散时间点的状态值。

分为前向欧拉和后向欧拉法（左图前向，右图后向）:

梯形法（Trapezoidal Method）

梯形法是一种更精确的数值积分方法，基于梯形规则对连续时间系统进行离散化。
它通过将积分区间内的变化近似为线性变化，从而提高了离散化的精度。
对于连续时间微分方程 $\frac{dx(t)}{dt} = f(x(t),t)$ ，梯形法的离散形式为： $x_{k+1}=x_k+ \frac{\triangle t}{2}(f(x_k,t_k)+f(x_{k+1},t_{k+1}))$

零阶保持法（Zero-Order Hold, ZOH）

零阶保持法是一种常用的离散化方法，主要用于将连续时间系统转换为离散时间系统。
它假设在采样间隔内，系统的输入保持不变（即零阶保持），从而将连续时间系统离散化。
在控制系统和信号处理中，ZOH 常用于将连续时间传递函数转换为离散时间传递函数。

从SSM到S4_state space model-CSDN博客 1.7章节有系统离散化推导过程

在 A和 B离散化后，使用步长 Δ 的离散版本可以改写为：

最后，模型通过全局卷积计算输出：

其中 M 是输入序列 x的长度， $K\in \mathbb{R}^M$ 是一个结构化卷积核。

SSM的卷积形式推导

卷积形式只有一个表达式： $y_k=x_k * kernel$

从SSM到S4_state space model-CSDN博客 2.2章节有SSM 的卷积表示详细推导过程

3.2 Vision Mamba概述

图释：我们首先将输入图像分割成多个 patch，然后将这些 patch 投影为 patch token。最后，我们将 token 序列输入到所提出的 Vim 编码器中。为了执行 ImageNet 分类任务，我们在 patch token 序列中额外连接一个可学习的分类 token。与用于文本序列建模的 Mamba 不同，Vim 编码器同时以正向和反向两个方向处理 token 序列。

图 2 展示了 Vim 的总体设计。标准的 Mamba 是为一维序列设计的。

为了处理视觉任务，我们首先将二维图像 $t\in \mathbb{R}^{H\times W\times C}$ 转换为扁平化的二维 patch 序列，其中 (H,W)是输入图像的尺寸，C 是通道数，P 是 patch 的大小。

接下来，我们将 $x_p$ 线性投影到维度为 D 的向量，并添加位置嵌入，如下所示：

其中 $t_j$ 是图像 t 的第 j 个 patch，是可学习的投影矩阵。受 ViT（Dosovitskiy et al., 2020）和 BERT（Kenton & Toutanova, 2019）启发，我们也使用 class token 表示整个 patch 序列，记为 $t_{cls}$ 。

然后，我们将 token 序列 $T_{l-1}$ 输入到 Vim 编码器的第 $l$ 层，得到输出 $T_l$ 。

最后，我们对输出的class token ${T_0}^L$ 进行归一化，并将其输入到多层感知机（MLP）头中以获得最终预测 $\hat{p}$ ，如下所示：

,。

其中 Vim 是提出的 Vision Mamba 模块，L 是层数，Norm 是归一化层。

3.3 Vim 模块

原始的 Mamba 模块是为一维序列设计的，不适合需要空间感知的视觉任务。本节介绍 Vim 模块，它结合了双向序列建模以适应视觉任务。图 2 展示了 Vim 模块的设计。

具体来说，我们在算法 1 中描述了 Vim 模块的操作。

输入 token 序列 $T_{l-1}$ 首先通过归一化层进行归一化。

接下来，我们将归一化后的序列线性投影到维度为 E的 x 和 z。

然后，我们从正向和反向两个方向处理 x。对于每个方向，我们首先对 x 应用一维卷积得到 ${x}'_o$ ，然后分别将 ${x}'_o$ 线性投影到 $B_o,C_o,\triangle _o$ 。接着，我们使用 $\triangle _o$ 对 $\bar{A}_o$ , $\bar{B}_o$ 进行转换。

最后，我们通过 SSM 计算 $y_{forward}$ 和 $y_{backward}$ ，并通过 z 进行门控后相加，得到输出 token 序列 $T_l$ 。

3.4 架构细节

总结来说，我们的架构超参数如下：

L 表示模块数量，D 表示隐藏状态维度，E 表示扩展状态维度，N 表示 SSM 维度。

受 ViT和 DeiT启发，我们首先使用 16×16 的核大小投影层获取非重叠 patch 嵌入的一维序列。

随后，我们直接堆叠 L 个 Vim 模块。默认情况下，我们将模块数量 L设置为 24，SSM 维度 N 设置为 16。为了与 DeiT 系列的模型大小对齐，对于 tiny 版本，我们将隐藏状态维度 D 设置为 192，扩展状态维度 E 设置为 384；对于 small 版本，我们将 D 设置为 384，E 设置为 768。

3.5 效率分析

传统的基于 SSM 的方法利用快速傅里叶变换（FFT）来加速卷积操作，如公式 (3) 所示。

对于数据依赖的方法（如 Mamba），算法 1 中第 11 行的 SSM 操作不能再使用卷积操作。？？

“对于数据依赖的方法（如 Mamba），算法 1 中第 11 行的 SSM 操作不能再使用卷积操作。”这句话什么意思？

答：这里说Mamba属于数据依赖方法是因为Mamba在SSM的基础上，为了使得模型具有过滤无用信息的能力，将SSM公式中的B,C矩阵由输入动态生成。但这样就会导致不能使用卷积表示，实现并行计算（因为卷积必须使用固定的kernel）。

详细请看Mamba理论入门_mamba一次性完成前向传播-CSDN博客 3.1 章节

为了解决这个问题，Mamba 和 Vim 选择了一种现代硬件友好的方式来确保效率。

优化的核心思想是避免现代硬件加速器（如 GPU）的 IO 瓶颈和内存瓶颈。

这部分可详看Mamba理论入门_mamba一次性完成前向传播-CSDN博客 3.3章节硬件感知算法。

IO 效率

高带宽内存（HBM）和 SRAM 是 GPU 的两个重要组件。其中，SRAM 具有更大的带宽，而 HBM 具有更大的内存容量。标准实现中，Vim 的 SSM 操作使用 HBM 时，内存 IO 的数量为 $O(BMEN)$ 。

受 Mamba 启发，Vim 首先从 HBM 读取 $O(BME+EN)$ 字节的内存 $(\triangle _o,A_o,B_o,C_o)$ 到 SRAM。然后，Vim 在 SRAM 中获取大小为 $(B,M,E,N)$ 的离散 $A_o,B_o$ 。最后，Vim 在 SRAM 中执行 SSM 操作，并将大小为 $(B,M,E)$ 的输出写回 HBM。这种方法可以将 IO 从 $O(BMEN)$ 减少到 $O(BME+EN)$ 。

在没有使用硬件感知算法之前，频繁的数据在 DRAM/HBM（全局内存） 和 SRAM（GPU 的共享存储器或寄存器） 之间的读写操作，显著降低性能：

Mamba提出kernel Fusion通过将多个计算操作合并到一个 CUDA kernel 中，从而减少不必要的内存访问和 kernel 调用，提高性能：

内存效率

为了避免内存溢出问题并降低处理长序列时的内存使用，Vim 采用了与 Mamba 相同的重计算方法。

对于大小为 $(B,M,E,N)$ 的中间状态，Vim 在网络反向传播时重新计算它们。对于激活函数的输出和卷积的中间激活值，Vim 也重新计算它们以优化 GPU 内存需求，因为激活值占用大量内存但重计算速度较快。

重计算Recomputation

计算效率

Vim 模块中的 SSM（算法 1 第 11 行）和 Transformer 中的自注意力机制都起到了关键作用，即自适应地提供全局上下文。给定一个视觉序列和默认设置，全局自注意力和 SSM 的计算复杂度分别为：

其中自注意力的复杂度与序列长度 M 成二次关系，而 SSM 的复杂度与序列长度 M 成线性关系（N 是固定参数，默认设置为 16）。这种计算效率使得 Vim 在处理大序列长度（如千兆像素应用）时具有可扩展性。

四、实验

4.1 图像分类设置

设置

我们在 ImageNet-1K 数据集（Deng et al., 2009）上对 Vim 进行基准测试，该数据集包含 1.28M 张训练图像和 50K 张验证图像，涵盖 1,000 个类别。所有模型均在训练集上训练，并在验证集上报告 top-1 准确率。为了公平比较，我们的训练设置主要遵循 DeiT（Touvron et al., 2021b）。具体来说，我们采用随机裁剪、随机水平翻转、标签平滑正则化、mixup 和随机擦除作为数据增强方法。在训练 224×224 输入图像时，我们使用 AdamW（Loshchilov & Hutter, 2019）优化模型，动量为 0.9，总批量大小为 1024，权重衰减为 0.05。我们使用余弦调度、初始学习率 1×10⁻³ 和指数移动平均（EMA）对 Vim 模型进行 300 轮训练。在测试阶段，我们对验证集应用中心裁剪以生成 224×224 图像。实验在 8 块 A800 GPU 上进行。

长序列微调

为了充分利用 Vim 的长序列建模能力，我们在预训练后继续以长序列设置对 Vim 进行 30 轮微调。具体来说，我们将 patch 提取步长设置为 8，同时保持 patch 大小不变，学习率恒定为 10⁻⁵，权重衰减为 10⁻⁸。

结果

表 1 比较了 Vim 与基于卷积网络（ConvNet）、Transformer 和 SSM 的骨干网络。与基于卷积网络的 ResNet（He et al., 2016）相比，Vim 表现出更优的性能。

例如，在参数量大致相同的情况下，VimSmall 的 top-1 准确率达到 80.3，比 ResNet50 高出 4.1 个百分点。与传统的基于自注意力的 ViT（Dosovitskiy et al., 2020）相比，Vim 在参数量和分类准确率上均显著优于 ViT。

与高度优化的 ViT 变体 DeiT（Touvron et al., 2021b）相比，Vim 在不同规模下均超越了 DeiT，且参数量相当：Vim-Tiny 比 DeiT-Tiny 高出 3.9 个百分点，Vim-Small 比 DeiT-Small 高出 0.5 个百分点，Vim-Base 比 DeiT-Base 高出 0.1 个百分点。

与基于 SSM 的 S4ND-ViT-B（Nguyen et al., 2022）相比，Vim 在参数量减少 3 倍的情况下实现了相似的 top-1 准确率。

经过长序列微调后，Vim-Tiny†、Vim-S† 和 Vim-B† 均取得了更高的结果，其中 Vim-S† 甚至达到了与 DeiT-B 相似的结果。

这些结果表明，Vim 能够轻松适应长序列建模，并提取更强的视觉表示。图 1 (b) 和 (c) 比较了 tiny 版本的 Vim 和 DeiT 的 FPS 和 GPU 内存。随着图像分辨率的增加，Vim 在速度和内存方面表现出更高的效率。具体来说，当图像大小为 512×512 时，Vim 的 FPS 和内存与 DeiT 相似；当图像大小增加到 1248×1248 时，Vim 的速度比 DeiT 快 2.8 倍，并节省了 86.8% 的 GPU 内存。Vim 在线性序列长度扩展方面的显著优势使其能够胜任高分辨率下游视觉应用和长序列多模态应用。

4.2 语义分割设置

设置

我们在 ADE20K（Zhou et al., 2019）数据集上进行语义分割实验，并使用 UperNet（Xiao et al., 2018b）作为分割框架。具体设置详见附录 B。

附录B

我们在 ADE20K（Zhou et al., 2019）数据集上进行语义分割实验。ADE20K 包含 150 个细粒度语义类别，分别有 20K、2K 和 3K 张图像用于训练、验证和测试。我们选择 UperNet（Xiao et al., 2018a）作为基础框架。在训练中，我们使用 AdamW 优化模型，权重衰减为 0.01，总批量大小为 16。采用的训练调度包括初始学习率 6×10⁻⁵、线性学习率衰减、1,500 次迭代的线性预热和总计 160K 次迭代的训练。数据增强遵循常见设置，包括随机水平翻转、在比例范围 [0.5, 2.0] 内的随机重新缩放以及随机光度失真。在评估期间，我们将图像缩放到较短边为 512。

结果

如表 2 所示，Vim 在不同规模下均优于 DeiT：Vim-Ti 比 DeiT-Ti 高出 1.8 mIoU，Vim-S 比 DeiT-S 高出 0.9 mIoU。

与 ResNet-101 骨干网络相比，我们的 Vim-S 在参数量减少近 2 倍的情况下实现了相同的分割性能。

为了进一步评估下游任务（如分割、检测和实例分割）的效率，我们将骨干网络与常用的特征金字塔网络（FPN）模块结合，并对其 FPS 和 GPU 内存进行基准测试。

如图 3 和图 4 所示，尽管我们在骨干网络上附加了较重的 FPN，效率曲线仍显示出与纯骨干网络（图 1）相似的比较结果。

这种优异的线性扩展性能归功于我们提出的高效骨干网络 Vim，它为学习千兆像素级视觉表示奠定了基础，无需多阶段编码（例如航空图像、医学图像和计算病理学）。

4.3 目标检测和实例分割设置

设置

我们在 COCO 2017 数据集（Lin et al., 2014）上进行目标检测和实例分割实验，并使用 ViTDet（Xiao et al., 2018b）作为基本框架。具体设置详见附录 B。

附录B

我们在 COCO 2017 数据集（Lin et al., 2014）上进行目标检测和实例分割实验。COCO 2017 数据集包含 118K 张训练图像、5K 张验证图像和 20K 张测试图像。我们使用经典的 Cascade Mask R-CNN（Cai & Vasconcelos, 2019）作为基础框架。对于基于 ViT 的骨干网络，我们应用额外配置（例如交错窗口和全局注意力）以处理高分辨率图像，遵循 ViTDet（Li et al., 2022c）。对于基于 SSM 的 Vim，我们直接使用它而无需任何修改。其他训练和评估设置相同。在训练中，我们使用 AdamW 优化模型，权重衰减为 0.1，总批量大小为 64。采用的训练调度包括初始学习率 1×10⁻⁴、线性学习率衰减和总计 380K 次迭代的训练。数据增强使用大规模抖动数据增强（Ghiasi et al., 2021）将输入图像调整为 1024×1024。在评估期间，我们将图像缩放到较短边为 1024。

结果

表 3 比较了 Vim-Ti 和 DeiT-Ti 在 Cascade Mask R-CNN 框架（Cai & Vasconcelos, 2019）下的表现。Vim-Ti 在 box AP 和 mask AP 上分别比 DeiT-Ti 高出 1.3 和 1.1。

对于中等大小和大型目标，Vim-Ti 在 box AP 和 mask AP 上分别比 DeiT-Ti 高出 1.6/1.3 和 1.4/1.8，显示出比 DeiT 更好的长距离上下文学习能力（图 5）。

我们强调，这种准确性优势并非微不足道，因为 DeiT 配备了窗口注意力机制，而 Vim 则采用纯序列建模方式。具体来说，为了在高分辨率图像（如 1024×1024）上进行表示学习，我们遵循 ViTDet（Li et al., 2022c）并修改了 DeiT 骨干网络，使用 2D 窗口注意力机制，这引入了 2D 先验并打破了 Transformer 的序列建模特性。得益于第 3.5 节、图 1 和图 4 中展示的效率，我们可以直接在 1024×1024 输入图像上应用 Vim，并学习序列视觉表示，而无需在骨干网络中引入 2D 先验。

4.4 消融实验

双向 SSM

我们消融了 Vim 的关键双向设计，使用 ImageNet-1K 分类和 Segmenter（Strudel et al., 2021）语义分割框架在 ADE20K 上进行实验。为了全面评估 ImageNet 上学习到的表示能力，我们使用仅包含 2 层的简单 Segmenter 头进行语义分割的迁移学习。我们研究了以下双向策略：

无：直接采用 Mamba 模块处理视觉序列，仅使用正向方向。
双向序列：在训练期间随机翻转视觉序列，类似于数据增强。
双向块：将堆叠的模块配对，每对的第一个模块处理正向序列，第二个模块处理反向序列。
双向 SSM：为每个模块添加额外的 SSM 以处理反向序列。
双向 SSM + Conv1d：在双向 SSM 的基础上，进一步在反向 SSM 前添加反向 Conv1d（图 2）。

如表 4 所示，直接采用 Mamba 模块在分类任务中表现良好，但单向方式在下游密集预测任务中面临挑战。具体来说，使用双向块的初步策略在分类任务中的 top-1 准确率降低了 7 个百分点，但在语义分割任务中比单向 Mamba 模块高出 1.3 mIoU。通过添加额外的反向 SSM 和 Conv1d，我们在分类准确率（73.9 top-1 acc vs. 73.2 top-1 acc）和分割性能（35.9 mIoU vs. 32.3 mIoU）上均取得了显著提升。我们将双向 SSM + Conv1d 策略作为 Vim 模块的默认设置。

分类设计

我们消融了 Vim 的分类设计，在 ImageNet-1K 分类任务上进行基准测试。我们研究了以下分类策略：

均值池化：对最后一个 Vim 模块的输出特征进行均值池化，并对池化后的特征进行分类。
最大池化：首先对视觉序列的每个 token 进行分类头适配，然后对序列进行最大池化以获取分类预测结果。
头部 class token：遵循 DeiT（Touvron et al., 2021b），在视觉序列的头部连接 class token 并进行分类。
双 class token：在头部 class token 策略的基础上，额外在序列尾部添加一个 class token。
中间 class token：在序列中间添加一个 class token，并对最终的中间 class token 进行分类。

如表 5 所示，实验表明中间 class token 策略能够充分利用 SSM 的循环特性和 ImageNet 中的中心对象先验，取得了最佳的 top-1 准确率 76.1。

五、结论与未来工作

我们提出了 Vision Mamba (Vim)，以探索最新的高效状态空间模型（即 Mamba）作为通用视觉骨干网络。与之前用于视觉任务的状态空间模型（使用混合架构或等效的全局 2D 卷积核）不同，Vim 以序列建模的方式学习视觉表示，并且没有引入图像特定的归纳偏置。得益于提出的双向状态空间建模，Vim 实现了数据依赖的全局视觉上下文，并具有与 Transformer 相同的建模能力，同时计算复杂度更低。受益于 Mamba 的硬件感知设计，Vim 在处理高分辨率图像时的推理速度和内存使用显著优于 ViTs。在标准计算机视觉基准上的实验结果验证了 Vim 的建模能力和高效性，表明 Vim 有潜力成为下一代视觉骨干网络。

在未来的工作中，Vim 的双向 SSM 建模与位置嵌入的结合非常适合无监督任务，例如掩码图像建模预训练。与 Mamba 类似的架构也支持多模态任务，例如 CLIP 风格的预训练。基于预训练的 Vim 权重，探索 Vim 在高分辨率医学图像、遥感图像和长视频分析中的实用性（这些可以被视为下游任务）将非常直接。

附录C——分层架构的扩展比较

为了进一步与分层架构进行比较，我们提出了另一种变体 Hier-Vim，通过将 SwinTransformer 中的局部窗口注意力替换为提出的全局双向 SSM。我们在表 6 中详细说明了配置。

ImageNet 分类

遵循标准的训练和验证协议（Liu et al., 2021; 2024），我们在表 7 中比较了 Hier-Vim 与流行的分层架构在 tiny、small 和 base 模型规模下的表现。结果表明，Hier-Vim 在 tiny 规模上比 Swin Transformer 高出 1.3%，在 small 规模上高出 0.2%，在 base 规模上高出 0.4%，显示出与成熟且高度优化的现代分层架构相竞争的性能。