作者 :Yue Liu 1 、Yunjie Tian1 、Yuzhong Zhao1 、Hongtian Yu1、 Lingxi Xie2 、Yaowei Wang3 、Qixiang Ye1 、Jianbin Jiao1 、Yunfan Liu1‡
单位 :中国科学院大学、华为、鹏城实验室
摘要:
设计高效的计算网络架构一直是计算机视觉领域的一项重要课题。在本文中,我们将一种名为 Mamba 的状态空间语言模型适配为 VMamba,一种具备线性时间复杂度的视觉主干网络(vision backbone)。VMamba 的核心是由多个视觉状态空间(Visual State-Space, VSS)模块和二维选择性扫描(2D Selective Scan, SS2D)模块构成的堆叠结构。通过沿四种扫描路径遍历,SS2D 在一维选择性扫描的有序性与二维视觉数据的非序列性之间架起桥梁,从而便于从多种来源和角度收集上下文信息。基于 VSS 模块,我们开发了一系列 VMamba 架构,并通过一系列架构和实现上的优化加速了它们的性能。大量实验证明,VMamba 在多种视觉感知任务中表现出色,尤其在输入扩展效率方面,相较于现有基准模型展现了明显优势。源码已开源。
脑图
一、Introduction
视觉表示学习作为计算机视觉领域的基础研究方向之一,在深度学习时代取得了显著进展。为了表征视觉数据中的复杂模式,目前主要有两类主干网络被提出并广泛应用于各种视觉任务中,即卷积神经网络(Convolutional Neural Networks, CNNs)[49, 27, 29, 53, 37] 和视觉 Transformer(Vision Transformers, ViTs)[13, 36, 57, 66]。与 CNNs 相比,ViTs 通常因其引入了自注意力机制 [58, 13],在处理大规模数据时表现出更强的学习能力。然而,自注意力机制的计算复杂度与 token 数量呈二次增长,这在涉及大空间分辨率的下游任务中会带来巨大的计算开销。
为了解决这一问题,许多研究致力于提升注意力计算的效率 [54, 36, 12]。然而,现有方法要么限制了有效感受野的大小 [36],要么在各种任务中表现出显著的性能下降 [30, 60]。这促使我们开发一种适用于视觉数据的新型架构,同时保留原始自注意力机制的固有优势,即全局感受野和动态加权参数 [23]。
最近,在自然语言处理(NLP)领域,Mamba [17] 作为一种创新的状态空间模型(State Space Model, SSM)[17, 43, 59, 71, 48],展示了在线性复杂度下处理长序列建模的潜力。受到这一进展的启发,我们提出 VMamba,这是一种集成了基于 SSM 的模块的新型视觉主干网络,可实现高效的视觉表示学习。然而,Mamba 的核心算法,即并行化的选择性扫描操作,本质上是为处理一维序列数据设计的。当将其适配于视觉数据(视觉数据缺乏内在的序列结构)时,会面临显著的挑战。为了解决这一问题,我们提出了二维选择性扫描(2D Selective Scan, SS2D),这是一种专为空间域遍历设计的四向扫描机制。与自注意力机制(如图 1(a) 所示)相比,SS2D 确保每个图像 patch 仅通过其对应扫描路径中压缩的隐藏状态获取上下文知识(如图 1(b) 所示),从而将计算复杂度从二次降低为线性。
基于 VSS 模块,我们开发了一系列 VMamba 架构(即 VMamba-Tiny/Small/Base),并通过架构改进和实现优化进一步提升其性能。与基于 CNN(如 ConvNeXt [37])、ViT(如 Swin [36], HiViT [66])和 SSM(如 S4ND [44], Vim [69])的基准视觉模型相比,VMamba 在 ImageNet-1K [9] 数据集上的不同模型规模下均实现了更高的图像分类精度。例如,VMamba-Base 的 top-1 精度达到 83.9%,比 Swin 高出 0.4%,且吞吐率比 Swin 高出超过 40%(646 vs. 458)。VMamba 的优势还体现在多个下游任务中,其中 VMamba-Tiny/Small/Base 在 COCO [33] 数据集上的目标检测任务( 1 × 1 \times 1× 训练策略)中分别实现了 47.3% / 48.7% / 49.2% 的 mAP,相较于 Swin 提高了 4.6% / 3.9% / 2.3%,相较于 ConvNeXt 提高了 3.1% / 3.3% / 2.2%。在 ADE20K [68] 数据集上的单尺度语义分割任务中,VMamba-Tiny/Small/Base 分别实现了 47.9% / 50.6% / 51.0% 的 mIoU,相较于 Swin 提高了 3.4% / 3.0% / 2.9%,相较于 ConvNeXt 提高了 1.9% / 1.9% / 1.9%。此外,与计算复杂度随着输入 token 数量呈二次增长的 ViT 模型不同,VMamba 在保持性能相当的同时,计算复杂度仅呈线性增长,展示了其在输入扩展方面的最先进表现。
二、相关工作
卷积神经网络(CNNs)
自从AlexNet [31] 提出以来,大量研究致力于提升基于CNN的模型在各种视觉任务中的建模能力 [49, 52, 27, 29] 和计算效率 [28, 53, 64, 46]。例如,深度卷积 [28] 和可变形卷积 [5, 70] 等复杂操作符的引入显著提高了CNN的灵活性和效能。最近,受Transformer成功应用的启发 [58],现代化的CNN模型 [37] 通过整合长程依赖 [11, 47, 34] 和动态权重 [23] 等机制展现了极具潜力的性能。
视觉Transformer(ViTs)
ViT [13] 作为先驱性工作,探讨了基于原生Transformer架构的视觉模型的有效性,并强调了大规模预训练在图像分类性能中的重要性。为减少ViT对大型数据集的依赖,DeiT [57] 提出了教师-学生蒸馏策略,将CNN的知识迁移至ViT,并凸显了归纳偏差(inductive bias)在视觉感知中的关键作用。在这一方向的启发下,随后的研究提出了分层结构的ViT [36, 12, 61, 39, 66, 55, 6, 10, 67, 1]。
另一个研究方向集中于提升自注意力机制的计算效率,后者是ViT的核心。Linear Attention [30] 通过将自注意力机制重新表述为核特征图的线性点积,利用矩阵乘法的结合性将复杂度从平方降低为线性。GLA [65] 提出了硬件友好的线性注意力变体,平衡了内存传输与并行化能力。RWKV [45] 利用线性注意力机制将Transformer的并行训练与RNN的高效推理相结合。RetNet [51] 引入了门控机制以实现并行化计算路径,为递归提供了替代方案。RMT [15] 则通过将时间衰减机制应用于空间领域,扩展了视觉表示学习。
状态空间模型(SSMs)
尽管ViT架构广泛应用于视觉任务,但由于自注意力机制的平方复杂度,尤其是在处理长输入序列(如高分辨率图像)时,仍面临显著挑战。在提升扩展效率的研究 [8, 7, 45, 51, 41] 中,SSM作为Transformer的有力替代方案,吸引了广泛关注。Gu等人 [21] 展示了基于HiPPO初始化 [18] 的SSM模型在处理长程依赖方面的潜力。为提高实际可行性,S4 [20] 提出了将参数矩阵归一化为对角结构的方法。自此,各种结构化SSM模型相继出现,包括复杂对角结构 [22, 19]、支持多输入多输出 [50]、对角加低秩分解 [24] 和选择机制 [17] 等架构增强技术。这些进展也被整合到更大规模的表示模型中 [43, 41, 16],进一步凸显了结构化状态空间模型在各种应用中的通用性和扩展性。
感悟:Baseline是Vision-Transformer,本文的灵感来自于Mamba,将输入层进行修改,做成选择性扫描输入,取代原有的输入方法。
三、预备知识
状态空间模型(SSMs)的数学形式:状态空间模型起源于卡尔曼滤波器 [32],它是一种线性时不变(LTI)的系统,可以通过隐藏状态
h
(
t
)
∈
R
N
\mathbf{h}(t) \in \mathbb{R}^{N}
h(t)∈RN 将输入信号
u
(
t
)
∈
R
u(t) \in \mathbb{R}
u(t)∈R 映射为输出响应
y
(
t
)
∈
R
y(t) \in \mathbb{R}
y(t)∈R。具体来说,连续时间状态空间模型可用线性常微分方程(ODE)表示如下:
其中,
A
∈
R
N
×
N
\mathbf{A} \in \mathbb{R}^{N \times N}
A∈RN×N、
B
∈
R
N
×
1
\mathbf{B} \in \mathbb{R}^{N \times 1}
B∈RN×1、
C
∈
R
1
×
N
\mathbf{C} \in \mathbb{R}^{1 \times N}
C∈R1×N 和
D
∈
R
1
D \in \mathbb{R}^{1}
D∈R1 是模型的权重参数。
状态空间模型的离散化 为了将连续时间状态空间模型(SSMs)集成到深度学习模型中,必须对其进行离散化。具体来说,在时间间隔 [ t a , t b ] \left[t_{a}, t_{b}\right] [ta,tb] 内,隐藏状态变量 h ( t ) \mathbf{h}(t) h(t) 在 t = t b t = t_{b} t=tb 时的解析解可以表示为:
h
(
t
b
)
=
e
A
(
t
b
−
t
a
)
h
(
t
a
)
+
e
A
(
t
b
−
t
a
)
∫
t
a
t
b
B
(
τ
)
u
(
τ
)
e
−
A
(
τ
−
t
a
)
d
τ
\mathbf{h}\left(t_b\right)=e^{\mathbf{A}\left(t_b-t_a\right)} \mathbf{h}\left(t_a\right)+e^{\mathbf{A}\left(t_b-t_a\right)} \int_{t_a}^{t_b} \mathbf{B}(\tau) u(\tau) e^{-\mathbf{A}\left(\tau-t_a\right)} d \tau
h(tb)=eA(tb−ta)h(ta)+eA(tb−ta)∫tatbB(τ)u(τ)e−A(τ−ta)dτ
通过以时间尺度参数
Δ
\Delta
Δ 进行采样(即
d
τ
∣
t
i
t
i
+
1
=
Δ
i
\left.d \tau \right|{t{i}}^{t_{i+1}} = \Delta_{i}
dτ∣titi+1=Δi),
h
(
t
b
)
\mathbf{h}(t_b)
h(tb) 可以离散化为:
h
b
=
e
A
(
Δ
a
+
…
+
Δ
b
−
1
)
(
h
a
+
∑
i
=
a
b
−
1
B
i
u
i
e
−
A
(
Δ
a
+
…
+
Δ
i
)
Δ
i
)
,
\mathbf{h}_b=e^{\mathbf{A}\left(\Delta_a+\ldots+\Delta_{b-1}\right)}\left(\mathbf{h}_a+\sum_{i=a}^{b-1} \mathbf{B}_i u_i e^{-\mathbf{A}\left(\Delta_a+\ldots+\Delta_i\right)} \Delta_i\right),
hb=eA(Δa+…+Δb−1)(ha+i=a∑b−1Biuie−A(Δa+…+Δi)Δi),
其中,
[
a
,
b
]
[a, b]
[a,b] 是对应的离散时间步区间。值得注意的是,该公式近似于零阶保持法(Zero-Order Hold, ZOH)方法的结果,而零阶保持法常用于状态空间模型相关研究(具体证明请参见附录 A)。
图2:2D选择性扫描(SS2D)的示意图。输入补丁沿着四种不同的扫描路径(交叉扫描)进行遍历,每个序列都由单独的S6块独立处理。然后将结果合并成一个二维特征映射作为最终输出(交叉合并)。
选择性扫描机制 为了应对线性时不变状态空间模型(LTI SSM)在捕获上下文信息方面的局限性,Gu 等人 [17] 提出了面向状态空间模型的全新参数化方法,该方法结合了一种基于输入选择的机制,称为 S6。然而,对于选择性状态空间模型,时间变化的权重参数带来了计算隐藏状态的效率挑战,因为卷积无法适应动态权重,从而难以直接应用。然而,由于公式 (3) 提供了隐藏状态 h b h_b hb 的递归关系,仍然可以利用关联扫描算法 [2, 42, 50] 高效地计算响应 y b y_b yb,其计算复杂度为线性(详细说明请参见附录 B)。
**思考与感悟:VMamba其实是采用四种顺序读取图片信息,然后相当于作为文本序列,放进Mamba里面进行训练。这里面要从代码去分析:
1.输入是否只是四种顺序去分块? 答:在VMamba/classification/models/csm_triton.py代码中有不同的Scan方式,
2.最后的四种结果是如何合并的? 答:按照分块的方法进行合并
3.每块是那种顺序分解? 答:分块之后并不是直接分解,而是逐元素乘上卷积核得到长度为D的embeding序列
4.代码是否调用了mamba的库? 答:使用的是ssd_minimal.py,这个代码不调用mamba的ssm库
第四章、VMamba
4.1网络结构
我们在三个规模下开发了VMamba:Tiny、Small和Base(分别称为VMamba-T、VMamba-S和VMamba-B)。VMamba-T的架构概览如图3(a)所示,详细配置见附录E。输入图像
I
∈
R
H
×
W
×
3
\mathbf{I} \in \mathbb{R}^{H \times W \times 3}
I∈RH×W×3首先通过一个stem模块分割为多个patch,得到一个空间维度为
H
/
4
×
W
/
4
H / 4 \times W / 4
H/4×W/4的2D特征图。没有引入额外的位置嵌入,多个网络阶段被用来创建具有
H
/
8
×
W
/
8
H / 8 \times W / 8
H/8×W/8、
H
/
16
×
W
/
16
H / 16 \times W / 16
H/16×W/16和
H
/
32
×
W
/
32
H / 32 \times W / 32
H/32×W/32分辨率的分层表示。具体来说,每个阶段包括一个下采样层(第一个阶段除外),然后是一个堆叠的视觉状态空间(VSS)块。
VSS块作为Mamba块[17](图3(b))的视觉对应物,用于表示学习。VSS块的初始架构(图3(c)中称为“香草VSS块”)通过替换S6模块来制定。S6是Mamba的核心,能够实现全局感受野、动态权重(即选择性)和线性复杂度。我们用新提出的2D-Selective-Scan(SS2D)模块替换它,更多细节将在接下来的子节中介绍。为了进一步提高计算效率,我们去除了整个乘法分支(如图3(c)中红框所示),因为门控机制的效果已通过SS2D的选择性得以实现。因此,改进后的VSS块(如图3(d)所示)由单个网络分支和两个残差模块组成,模拟了香草Transformer块[58]的架构。本文中的所有结果均使用基于该架构的VSS块构建的VMamba模型获得。
4.2 2D版本的选择状态(SS2D)
尽管S6模块中扫描操作的顺序性质非常适合处理涉及时间数据的NLP任务,但当应用于视觉数据时(视觉数据本质上是非顺序的,并且包含空间信息,如局部纹理和全局结构),却带来了显著挑战。为了解决这一问题,S4ND [44]通过卷积操作重新表述了SSM,直接通过外积将卷积核从1D扩展到2D。然而,这种修改限制了权重不依赖于输入,从而导致捕获上下文信息的能力受限。因此,我们坚持使用选择性扫描方法[17]处理输入,并提出了2D-Selective-Scan(SS2D)模块,以适应视觉数据,同时不妥协其优势。
图3:左:VMamba整体架构的示意图(a),以及Mamba和VSS块的结构图(b)到(d)。右:在分类精度和计算效率方面,VMamba变体与基准方法的比较。
图2展示了SS2D中的数据转发包括三个步骤:交叉扫描、使用S6块的选择性扫描以及交叉合并。具体来说,SS2D首先将输入的patch展平为沿四个不同遍历路径的序列(即交叉扫描)。然后,每个patch序列通过单独的S6块并行处理,最终将结果序列重新整形并合并形成输出图(即交叉合并)。通过使用互补的1D遍历路径,SS2D使得图像中的每个像素可以从不同方向上的所有其他像素处整合信息。这种整合有助于在2D空间中建立全局感受野。
4.3 加速VMamba
如图3(e)所示,使用香草VSS块的VMamba-T模型(称为“香草VMamba”)达到了426张图像/秒的吞吐量,并包含22.9M参数和5.6G FLOPs。尽管在Tiny级别上以82.2%的分类精度达到了最先进的水平(比Swin-T [36]高出0.9%),但低吞吐量和高内存开销对VMamba的实际部署提出了显著挑战。
在本小节中,我们概述了为提高推理速度所做的努力,主要集中在实现细节和架构设计的改进上。我们在ImageNet-1K上进行了图像分类模型评估。每个渐进式改进的影响总结如下,其中( % , i m g / s \%, \mathrm{img} / \mathrm{s} %,img/s)表示ImageNet-1K上的Top-1准确率和推理吞吐量的增益。进一步讨论见附录E。
步骤(a)(
+
0.0
%
,
+
41
i
m
g
/
s
+0.0\%, +41 \mathrm{img}/\mathrm{s}
+0.0%,+41img/s):通过在Triton中重新实现交叉扫描和交叉合并。
步骤(b)(
+
0.0
%
,
−
3
i
m
g
/
s
+0.0\%, -3 \mathrm{img}/\mathrm{s}
+0.0%,−3img/s):通过调整CUDA实现选择性扫描,以适应float16输入和float32输出。这显著提高了训练效率(吞吐量从165提升至184),尽管在测试时速度略有波动。
步骤(c)(
+
0.0
%
,
+
174
i
m
g
/
s
+0.0\%, +174 \mathrm{img}/\mathrm{s}
+0.0%,+174img/s):通过用线性变换(即torch.nn.functional.linear)替代选择性扫描中较慢的einsum。我们还采用了(
B
,
C
,
H
,
W
B, C, H, W
B,C,H,W)的张量布局,以消除不必要的数据排列。
表1:在ImageNet-1K上的性能比较。吞吐量值使用A100 GPU和AMD EPYC 7542 CPU进行测量,使用[62]发布的工具包,遵循[36]中提出的协议。所有图像的大小均为224×224。
步骤(d)(
−
0.6
%
,
+
175
i
m
g
/
s
-0.6\%, +175 \mathrm{img}/\mathrm{s}
−0.6%,+175img/s):由于其计算效率,引入MLP到VMamba中。我们还去除了DWConv(深度卷积[23])层,并将层配置从[2, 2, 9, 2]更改为[2, 2, 2, 2],以降低FLOPs。
步骤(e)(
+
0.6
%
,
+
366
i
m
g
/
s
+0.6\%, +366 \mathrm{img}/\mathrm{s}
+0.6%,+366img/s):通过将参数ssm-ratio(特征扩展因子)从2.0减少到1.0(也称为步骤(d.1)),将层数提高到[2, 2, 5, 2](也称为步骤(d.2)),并去除如图3(c)所示的整个乘法分支。
步骤(f)(
+
0.3
%
,
+
161
i
m
g
/
s
+0.3\%, +161 \mathrm{img}/\mathrm{s}
+0.3%,+161img/s):通过引入DWConv层(也称为步骤(e.1)),并将参数d_state(SSM状态维度)从16.0减少到1.0(也称为步骤(e.2)),同时将ssm-ratio提高回2.0。
步骤(g)(
+
0.1
%
,
+
346
i
m
g
/
s
+0.1\%, +346 \mathrm{img}/\mathrm{s}
+0.1%,+346img/s):通过将ssm-ratio减少到1.0,并将层配置从
[
2
,
2
,
5
,
2
]
[2, 2, 5, 2]
[2,2,5,2]更改为
[
2
,
2
,
8
,
2
]
[2, 2, 8, 2]
[2,2,8,2]。
第五章、实验
在本节中,我们展示了一系列实验来评估VMamba的性能,并与各种视觉任务中的流行基准模型进行比较。我们还通过与其他方法的比较,验证了所提的2D特征图遍历方法的有效性。此外,我们通过可视化其有效感受野(ERF)和激活图,分析了VMamba的特性,并考察了它在处理较长输入序列时的可扩展性。我们主要遵循了Swin模型[36]的超参数设置和实验配置。有关详细的实验设置,请参见附录E和F,更多的消融实验见附录H。所有实验均在配备 8 × 8 \times 8× NVIDIA Tesla-A100 GPU的服务器上进行。
5.1 图像分类
我们评估了VMamba在ImageNet-1K [9]上的图像分类性能,并将结果与基准方法进行了比较,结果总结在表1中。VMamba-T在相似的FLOPs下,达到了82.6%的top-1准确率,超越了DeiT-S 2.8%和Swin-T 1.3%。值得注意的是,VMamba在Small和Base尺度上也保持了性能优势。例如,VMamba-B达到了83.9%的top-1准确率,超越了DeiT-B 2.1%和Swin-B 0.4%。
在计算效率方面,VMamba-T实现了每秒1686张图像的吞吐量,优于或与最先进的方法相当。这一优势在VMamba-S和VMamba-B中继续体现,分别达到了每秒877张图像和646张图像。与基于SSM的模型相比,VMamba-T的吞吐量比S4ND-Conv-T [44]高出1.47倍,比Vim-S [69]高出1.08倍,同时在性能上分别领先这两种模型0.4%和2.1%。
表2:左侧:在MSCOCO上进行的目标检测和实例分割结果。 A P b AP^{b} APb和 A P m AP^{m} APm分别表示框AP和掩码AP。FLOPs是在输入尺寸为 1280 × 800 1280 \times 800 1280×800时计算的。’ 1 × 1 \times 1× ‘表示模型进行了12轮微调,’ 3 × M S 3 \times \mathrm{MS} 3×MS '表示进行了36轮的多尺度训练。右侧:在ADE20K上进行的语义分割结果。FLOPs是在输入尺寸为 512 × 2048 512 \times 2048 512×2048时计算的。'SS’和’MS’分别表示单尺度和多尺度测试。
5.2 下游任务
在本小节中,我们评估了VMamba在下游任务上的表现,包括在MSCOCO2017 [33]上的目标检测和实例分割,以及在ADE20K [68]上的语义分割。训练框架基于MMDetection [3]和MMSegmentation [4]库,采用[35]中的Mask R-CNN [26]和UperNet [63]作为检测和分割网络。
目标检测和实例分割。MSCOCO上的结果如表2所示。VMamba在不同训练计划下的框和掩码平均精度(AP b { }^{\mathrm{b}} b和 A P m AP^{\mathrm{m}} APm)上表现优越。在12轮微调计划下,VMamba-T/S/B的目标检测mAP分别为 47.3 % 47.3\% 47.3%、 48.7 % 48.7\% 48.7%、 49.2 % 49.2\% 49.2%,超越了Swin-T/S/B分别为 4.6 % 4.6\% 4.6%、 3.9 % 3.9\% 3.9%、 2.3 % 2.3\% 2.3% mAP,以及ConvNeXt-T/S/B分别为 3.1 % 3.1\% 3.1%、 3.3 % 3.3\% 3.3%、 2.2 % 2.2\% 2.2% mAP。VMamba-T/S/B的实例分割mAP超越了Swin-T/S/B分别为 3.4 % 3.4\% 3.4%、 2.8 % 2.8\% 2.8%、 1.8 % 1.8\% 1.8% mAP,超越了ConvNeXt-T/S/B分别为 2.6 % 2.6\% 2.6%、 1.9 % 1.9\% 1.9%、 1.4 % 1.4\% 1.4% mAP。此外,VMamba在使用36轮微调和多尺度训练时仍然保持优势,突显了其在需要密集预测的下游任务中的强大潜力。
语义分割。与之前的实验一致,VMamba在ADE20K上的语义分割性能优越,且具有相似的参数量。如表2所示,VMamba-T在单尺度(SS)设置下比Swin-T高出3.4%的mIoU,比ConvNeXt-T高出1.9%。在多尺度(MS)输入下,VMamba的优势依然持续。对于Small和Base尺度的模型,VMamba-S/B在SS设置下比NAT-S/B [25]高出
2.6
%
/
2.5
%
2.6\% / 2.5\%
2.6%/2.5% mIoU,在MS设置下高出
1.7
%
/
1.9
%
1.7\% / 1.9\%
1.7%/1.9% mIoU。
图4:VMamba对(a)下游任务和(b)以分辨率逐渐增加的输入图像的适应性说明。Swin-T∗表示使用缩放窗口大小进行的Swin-T测试。
图5:VMamba的资源消耗图。Swin-T∗表示使用缩放窗口大小测试的Swin-T
讨论。本小节中的实验结果展示了VMamba在目标检测、实例分割和语义分割任务中的适应性。在图4(a)中,我们将VMamba的表现与Swin和ConvNeXt进行了比较,突出显示了它在处理下游任务时,相比于ImageNet-1K上具有相似分类准确率的其他模型的优势。图4(b)进一步展示了VMamba在不同输入图像尺寸下的最稳定表现(即表现下降较小),在输入分辨率为 768 × 768 768 \times 768 768×768时,未进行微调的top-1分类准确率为 74.7 % 74.7\% 74.7%(线性调优后为 79.2 % 79.2\% 79.2%)。尽管输入分辨率增大,但VMamba展现出更大的容错性,同时FLOPs和内存消耗(见图5 (a)和©)呈线性增长,保持了高吞吐量(图5 (b)),使其在处理较大空间分辨率输入的下游任务时,比基于ViT的方法更有效和高效。这与Mamba在高效长序列建模方面的先进能力一致[17]。
5.3 分析
SS2D与自注意力的关系。为了在时间区间
[
a
,
b
]
[a, b]
[a,b]内计算响应
Y
Y
Y,我们将相应的SSM相关变量
u
i
⊙
Δ
i
∈
R
1
×
D
v
\mathbf{u}_{i} \odot \boldsymbol{\Delta}_{i} \in \mathbb{R}^{1 \times D_{v}}
ui⊙Δi∈R1×Dv,
B
i
∈
R
1
×
D
k
\mathbf{B}_{i} \in \mathbb{R}^{1 \times D_{k}}
Bi∈R1×Dk,和
C
i
∈
R
1
×
D
k
\mathbf{C}_{i} \in \mathbb{R}^{1 \times D_{k}}
Ci∈R1×Dk分别表示为
V
∈
R
T
×
D
v
\mathbf{V} \in \mathbb{R}^{T \times D_{v}}
V∈RT×Dv,
K
∈
R
T
×
D
k
\mathbf{K} \in \mathbb{R}^{T \times D_{k}}
K∈RT×Dk,和
Q
∈
R
T
×
D
k
\mathbf{Q} \in \mathbb{R}^{T \times D_{k}}
Q∈RT×Dk。因此,
y
b
\mathbf{y}_{\mathbf{b}}
yb沿
D
v
D_{v}
Dv维度的第
j
j
j个切片
y
b
(
j
)
∈
R
\mathbf{y}_{\mathbf{b}}^{(j)} \in \mathbb{R}
yb(j)∈R可以写成
y
b
(
j
)
=
(
Q
T
⊙
w
T
(
j
)
)
h
a
(
j
)
+
Q
T
∑
i
=
1
T
(
w
T
(
j
)
w
i
(
j
)
⊙
K
i
)
⊤
⊙
(
V
i
(
j
)
)
\mathbf{y}_{\mathbf{b}}^{(j)}=\left(\mathbf{Q}_{\mathbf{T}} \odot \mathbf{w}_{\mathbf{T}}^{(j)}\right) \mathbf{h}_{\mathbf{a}}^{(j)}+\mathbf{Q}_{\mathbf{T}} \sum_{i=1}^{T}\left(\frac{\mathbf{w}_{\mathbf{T}}^{(j)}}{\mathbf{w}_{\mathbf{i}}^{(j)}} \odot \mathbf{K}_{\mathbf{i}}\right)^{\top} \odot\left(\mathbf{V}_{\mathbf{i}}^{(j)}\right)
yb(j)=(QT⊙wT(j))ha(j)+QTi=1∑T(wi(j)wT(j)⊙Ki)⊤⊙(Vi(j))
图6:查询补丁的激活图,其中红星标记了查询补丁。通过组合SS2D中每个扫描路径的激活图(图6(b)和(c))得到的可视化结果。
图7:VMamba与其他基准模型之间的有效感受野(ERF)比较。响应较大的像素表示与中心像素相关的较大响应。
在公式(5)中,涉及 Q \mathbf{Q} Q、 K \mathbf{K} K和 V \mathbf{V} V的矩阵乘法过程与自注意力机制密切相似,尽管引入了 w \mathbf{w} w。
激活图的可视化。为了更直观和深入地理解SS2D,我们进一步可视化了对应于特定查询补丁的 Q K ⊤ \mathbf{QK}^{\top} QK⊤和 ( Q ⊙ w ) ( K / w ) ⊤ (\mathbf{Q} \odot \mathbf{w})(\mathbf{K} / \mathbf{w})^{\top} (Q⊙w)(K/w)⊤的注意力值(即激活图)。如图6(b)所示, Q K ⊤ \mathbf{QK}^{\top} QK⊤的激活图展示了SS2D在捕捉和保留遍历信息方面的有效性,所有先前扫描过的位于前景区域的token都被激活。图6(c)中包含了 w \mathbf{w} w的激活图,显示了更专注于查询补丁的邻域,符合 w \mathbf{w} w中固有的时间加权效果。扫描机制允许VMamba沿扫描路径积累历史信息,从而建立图像补丁间的长期依赖性。如图6(d)中红框内所示,较早扫描到的远左侧羊群的补丁仍然被激活。
有效感受野(ERF)的可视化。有效感受野(ERF)[40, 11]指的是输入空间中对特定输出单元激活有贡献的区域。我们对不同视觉骨干网络在训练前后的中心像素ERF进行了对比分析。图7展示的结果表明,在所有模型中,只有DeiT、HiViT、Vim和VMamba展示了全局ERF,而其他模型尽管理论上能够实现全局覆盖,但实际却展现了局部ERF。此外,VMamba的线性时间复杂度增强了其计算效率,相比于DeiT和HiViT在输入补丁数量上所付出的二次代价。尽管VMamba和Vim都基于Mamba架构,但VMamba的ERF比Vim的更均匀、2D-aware,这也可能直观地解释了它的优越性能。
选择性扫描模式的诊断研究。我们将提出的扫描模式(即Cross-Scan)与三种基准模式进行比较:单向扫描(Unidi-Scan)、双向扫描(Bidi-Scan)和级联扫描(Cascade-Scan,即按行和列依次扫描数据)。为了公平比较,我们调整了特征维度以保持相似的模型参数和FLOPs。
图8:不同扫描模式的性能比较。所提出的Cross-Scan在速度上优于其他扫描模式,同时保持了相同数量的参数和FLOPs。
第六章、结论
本文提出了VMamba,一种高效的视觉骨干网络模型,基于状态空间模型(SSM)。VMamba将自然语言处理任务中的选择性SSM优势引入到视觉数据处理中,弥合了有序的一维扫描和非顺序二维遍历之间的差距,采用了新颖的SS2D模块。此外,通过一系列的架构和实现优化,我们显著提高了VMamba的推理速度。通过广泛的实验,验证了VMamba家族的有效性,其线性时间复杂度使得VMamba在处理大分辨率输入的下游任务中具有优势。
局限性。尽管VMamba展示了有前景的实验结果,但本研究仍有改进空间。此前的研究已验证了在大规模数据集(如ImageNet-21K)上进行无监督预训练的有效性。然而,现有的预训练方法与基于SSM架构的VMamba兼容性,以及针对这类模型的预训练技术的研究,仍然未得到充分探索。探讨这些方面可能为未来的架构设计研究提供一个有前景的方向。此外,有限的计算资源阻碍了我们对VMamba在大规模上的架构探索,并进一步细化超参数搜索以提升实验性能。虽然SS2D是VMamba的核心组件,并且不对输入数据的布局或模态做出具体假设,使其可以广泛应用于多种任务,但VMamba在更通用任务中的潜力仍未得到充分发掘。弥合SS2D与这些任务之间的差距,并提出一种更加通用的视觉任务扫描模式,是一个有前景的研究方向。