MambaVision：一种Mamba-Transformer混合视觉骨干网络

摘要

我们提出了一种新型混合Mamba-Transformer主干网络，称为MambaVision，该网络专为视觉应用而设计。我们的核心贡献包括重新设计Mamba公式，以增强其对视觉特征的高效建模能力。此外，我们还对将视觉Transformer（ViT）与Mamba集成的可行性进行了全面的消融研究。我们的结果表明，在最终层添加几个自注意力模块可以显著提高捕捉长距离空间依赖关系的建模能力。基于我们的发现，我们引入了一系列具有分层架构的MambaVision模型，以满足各种设计标准。在ImageNet-1K数据集的图像分类任务中，MambaVision模型变体在Top-1准确率和图像吞吐量方面达到了新的最先进（SOTA）性能。在MS COCO和ADE20K数据集上的目标检测、实例分割和语义分割等下游任务中，MambaVision优于同等规模的骨干网络，并表现出更优越的性能。代码链接：https://github.com/NVlabs/MambaVision。

1 引言

近年来，Transformer[1]已成为包括计算机视觉、自然语言处理、语音处理和机器人技术在内的不同领域中的事实标准架构。此外，Transformer架构的通用性（主要归因于其注意力机制）及其灵活性，使其非常适合多模态学习任务，在这些任务中，整合和处理来自各种模态的信息至关重要。尽管有这些优点，但注意力机制相对于序列长度的二次复杂度使得Transformer在训练和部署时计算成本高昂。最近，Mamba[2]提出了一种新的状态空间模型（SSM），该模型实现了线性时间复杂度，并在不同的语言建模任务中优于或匹敌Transformer[2]。Mamba的核心贡献是一种新颖的选择机制，该机制能够实现对长序列的高效输入依赖处理，并考虑了硬件感知因素。

最近，还提出了一些基于Mamba的骨干网络[3, 4]，以利用其在图像分类和语义分割等视觉任务中的SSM公式的优势。然而，Mamba的自回归公式虽然对需要顺序数据处理的任务有效，但在受益于全感受野的计算机视觉任务中面临限制：（1）与顺序重要的序列不同，图像像素没有相同的顺序依赖关系。相反，空间关系通常是局部的，需要以更并行和集成的方式考虑。因此，这导致处理空间数据效率低下。（2）像Mamba这样的自回归模型逐步处理数据，限制了其在一个前向传递中捕捉和利用全局上下文的能力。相比之下，视觉任务通常需要理解全局上下文才能对局部区域进行准确预测。

Vision Mamba（Vim）[3]和其他研究提出了诸如双向SSM等修改，以解决缺乏全局上下文和空间理解的问题。虽然双向SSM有可能捕捉更全面的上下文，但由于需要在做出预测之前处理整个序列，因此会引入显著的延迟。此外，增加的复杂性可能导致训练挑战、过拟合风险，并且并不总是能提高准确性。由于这些缺陷，基于视觉Transformer（ViT）和卷积神经网络（CNN）架构的骨干网络仍然在不同视觉任务上优于最佳的基于Mamba的视觉模型。

在这项工作中，我们系统地重新设计了Mamba模块，使其更适合视觉任务。我们提出了一种混合架构，该架构由我们提出的公式（即MambaVision Mixer和MLP）以及Transformer模块组成。具体来说，我们研究了不同的集成模式，例如以等参数方式将Transformer模块添加到早期、中期和最终层以及每$l$层。我们的分析表明，在最终阶段利用几个自注意力模块可以显著增强捕捉全局上下文和长距离空间依赖关系的能力。如第5节所示，与纯Mamba或基于ViT的模型相比，使用混合架构还可以提高图像吞吐量。

我们引入了MambaVision模型，该模型由多分辨率架构组成，并利用基于CNN的残差模块进行大分辨率特征的快速特征提取。如图1所示，MambaVision在ImageNet-1K的Top-1准确率和图像吞吐量方面达到了新的SOTA帕累托前沿，优于Mamba、CNN和基于ViT的模型，有时优势显著。在目标检测和实例分割以及语义分割等下游任务中，具有MambaVision骨干网络的模型在MS COCO和ADE20数据集上分别优于同等规模的对应模型。因此，它验证了MambaVision作为高效骨干网络的有效性和通用性。

据我们所知，MambaVision是首次研究和开发包含Mamba和Transformer的混合架构以用于计算机视觉应用的工作。我们在这项工作中的主要贡献总结如下：

• 我们引入了一个重新设计的视觉友好型Mamba模块，提高了原始Mamba架构的准确率和图像吞吐量。
• 我们对Mamba和Transformer模块的集成模式进行了系统研究，并证明在最终阶段加入自注意力模块可以显著提高模型捕捉全局上下文和长距离空间依赖关系的能力。
• 我们引入了MambaVision，这是一种新型的混合Mamba Transformer模型。分层MambaVision在ImageNet-1K数据集上实现了Top-1准确率和图像吞吐量权衡方面的新SOTA帕累托前沿。

2 相关工作

ViT。视觉Transformer（ViT）[5]作为卷积神经网络（CNN）的一种有前景的替代方案出现，它利用自注意力层来提供更大的感受野。然而，ViT最初缺乏CNN固有的一些优势，如归纳偏置和平移不变性，并且需要大规模的训练数据集才能实现有竞争力的性能。为了解决这些限制，数据高效图像Transformer（DeiT）[6]引入了一种基于蒸馏的训练策略，显著提高了分类准确率，即使在较小的数据集上也是如此。在此基础上，LeViT[7]模型提出了一种混合方法，结合了为快速推理优化的重新设计的多层感知机（MLP）和自注意力模块，提高了效率和性能。此外，交叉协方差图像Transformer（XCiT）[8]引入了一种转置自注意力机制，有效地建模了特征通道之间的交互，提高了模型捕捉数据中复杂模式的能力。金字塔视觉Transformer（PVT）[9]采用了一种分层结构，在每个阶段的开始进行块嵌入和空间维度缩减，从而提高了计算效率。类似地，Swin Transformer[10]提出了一种分层架构，其中自注意力在局部窗口内计算，并通过移动窗口来实现区域之间的交互，平衡了局部和全局上下文。Twins Transformer[11]具有空间可分离自注意力，显著提高了效率。此外，Focal Transformer[12]利用焦点自注意力来捕捉长距离空间交互的细粒度细节。

Mamba。自Mamba引入以来，已经提出了许多努力来利用其能力进行视觉应用。具体来说，Vim[3]提出使用双向状态空间模型（SSM）公式，与Mamba公式相同，其中令牌在前向和后向方向上处理，以捕捉更多的全局上下文并改善空间理解。然而，双向编码增加了计算负载，可能会减慢训练和推理时间。此外，有效地结合来自多个方向的信息以形成连贯的全局理解是具有挑战性的，因为一些全局上下文可能会在过程中丢失。与Vim不同，我们提出的MambaVision使用单个前向传递和重新设计的Mamba块，可以捕捉短距离和长距离信息，并在ImageNet Top-1准确率和吞吐量方面显著优于它。

EfficientVMamba[4]提出了一种基于空洞的选择性扫描和跳跃采样方法，以有效地提取全局空间依赖关系。EfficientVMamba还使用了一种分层架构，由SSM和基于CNN的块组成，其中SSM用于较大的输入分辨率以更好地捕捉全局上下文，而CNN用于较低的分辨率。相比之下，MambaVision在较高分辨率下使用CNN进行快速特征提取，同时在较低分辨率下使用SSM和自注意力来捕捉短距离和长距离空间依赖关系的细粒度细节。我们提出的MambaVision在Top-1准确率和图像吞吐量方面也显著优于EfficientVMamba。

此外，VMamba[13]引入了一种基于Mamba的通用视觉骨干网络，具有交叉扫描模块（CSM），该模块促进了一维选择性扫描并扩大了全局感受野。具体来说，CSM模块采用四路选择性扫描方法（即左上和右下到相反方向）来整合来自所有周围令牌的信息并捕捉更多的全局上下文。此外，VMamba进行了架构更改，例如使用深度卷积和分层多分辨率结构。尽管CSM模块的设计更适合视觉任务，但感受野仍然受到交叉扫描路径的限制。与VMamba相比，我们提出的MambaVision混合器的设计更简单，可以捕捉短距离和长距离依赖关系。MambaVision还使用基于CNN的层进行快速特征提取，而VMamba在所有阶段都使用相同的块结构。此外，MambaVision模型在具有显著更高的图像吞吐量的同时优于VMamba对应模型。

3 方法

3.1 宏观架构

在本节中，我们介绍了MambaVision，这是我们提出的在ImageNet-1K数据集上具有最先进（SOTA）性能的新型架构。如图2所示，MambaVision具有分层架构，由4个不同的阶段组成。前两个阶段由基于CNN的层组成，用于在较高的输入分辨率下快速提取特征，而阶段3和4包括提出的MambaVision和Transformer块。具体来说，给定大小为$H\times W\times 3$的图像，输入首先被转换为大小为$\frac{H}{4}\times \frac{W}{4}\times C$的重叠块，并通过由两个连续的$3\times 3$ CNN层（步长为2）组成的stem投影到$C$维嵌入空间中。阶段之间的下采样器由一个批归一化的$3\times 3$ CNN层（步长为2）组成，该层将图像分辨率减半。此外，阶段1和2中的CNN块遵循通用的残差块公式，如下所示：
$\begin{array}{l}\hat{\mathbf{z}}=\mathrm{GELU}\left(\mathrm{BN}\left({\mathrm{Conv}}_{3\times 3}(\mathbf{z})\right)\right)\\ \mathbf{z}=\mathrm{BN}\left({\mathrm{Conv}}_{3\times 3}(\hat{\mathbf{z}})\right)+\mathbf{z}\end{array}$
其中，$\mathrm{GELU}$和$\mathrm{BN}$分别表示高斯误差线性单元激活函数和批归一化。

3.2 微观架构

在本节中，我们首先回顾Mamba和SSM的初步知识。然后，我们介绍阶段3和4的微观架构设计，并更详细地讨论MambaVision公式。

3.2.1 Mamba初步知识

Mamba是结构化状态空间序列模型（S4）的扩展，能够将一维连续输入$x(t)\in \mathbb{R}$转换为$y(t)\in \mathbb{R}$，通过可学习的隐藏状态$h(t)\in {\mathbb{R}}^M$和参数$\mathbf{A}\in {\mathbb{R}}^{M\times M}$、$\mathbf{B}\in {\mathbb{R}}^{1\times M}$和$\mathbf{C}\in {\mathbb{R}}^{1\times M}$，如下所示：
$$\begin{array}{rl}{h}^{\prime }(t)& =\mathbf{A}h(t)+\mathbf{B}x(t),\\ y(t)& =\mathbf{C}h(t)\end{array}$$

离散化 上述公式中的连续参数$\mathbf{A}$、$\mathbf{B}$和$\mathbf{C}$进一步转换为离散参数，以提高计算效率[16]。具体来说，假设一个时间尺度$\Delta$，可以应用零阶保持规则来获得离散参数$\overline{\mathbf{A}}\in {\mathbb{R}}^{M\times M}$、$\overline{\mathbf{B}}\in {\mathbb{R}}^{1\times M}$和$\overline{\mathbf{C}}\in {\mathbb{R}}^{1\times M}$，如下所示：
$$\begin{array}{l}\overline{\mathbf{A}}=\exp (\Delta \mathbf{A})\\ \overline{\mathbf{B}}=(\Delta \mathbf{A}{)}^{-1}(\exp (\Delta \mathbf{A})-\mathbf{I})\cdot (\Delta \mathbf{B}),\\ \overline{\mathbf{C}}=\mathbf{C}\end{array}$$
然后，公式（2）可以用离散参数表示为：
$$\begin{array}{rl}h(t)& =\overline{\mathbf{A}}h(t-1)+\overline{\mathbf{B}}x(t),\\ y(t)& =\overline{\mathbf{C}}h(t)\end{array}$$

此外，对于大小为$T$的输入序列，可以应用核为$\overline{\mathbf{K}}$的全局卷积来计算公式（4）的输出，如下所示：
$\begin{array}{rl}\overline{\mathbf{K}}& =\left(\mathbf{C}\overline{\mathbf{B}},\mathbf{C}\overline{\mathbf{AB}},\dots ,\mathbf{C}{\overline{\mathbf{A}}}^{T-1}\overline{\mathbf{B}}\right)\\ \mathbf{y}& =\mathbf{x}\ast \overline{\mathbf{K}}\end{array}$

选择性 Mamba通过引入选择机制进一步扩展了S4公式，该机制允许进行输入依赖的序列处理。这允许模型的参数$\mathbf{B}$、$\mathbf{C}$和$\Delta$根据输入动态调整，并过滤掉无关信息。进一步的离散化细节在[2]中提供。

3.2.2 层架构

假设输入$X\in {\mathbb{R}}^{T\times C}$具有序列长度$T$和嵌入维度$C$，阶段3和4中层$n$的输出可以计算为：
$$\begin{array}{l}{\hat{X}}^n=\mathrm{Mixer}\left(\mathrm{Norm}\left(X^{n-1}\right)\right)+X^{n-1}\\ X^n=\mathrm{MLP}\left(\mathrm{Norm}\left({\hat{X}}^n\right)\right)+{\hat{X}}^n\end{array}$$
其中，$\mathrm{Norm}$和$\mathrm{Mixer}$分别表示层归一化和令牌混合块的选择。不失一般性，使用层归一化作为$\mathrm{Norm}$。给定$N$层，前$\frac{N}{2}$层使用MambaVision混合器块，而剩余的$\frac{N}{2}$层使用自注意力。我们在下面描述每个混合器的细节。

MambaVision混合器 如图3所示，我们重新设计了原始的Mamba混合器，使其更适合视觉任务。首先，我们提出用常规卷积替换因果卷积，因为因果卷积将影响限制在一个方向上，这对视觉任务来说是不必要且有限制的。此外，我们添加了一个没有SSM的对称分支，由一个额外的卷积和SiLU激活组成，以补偿由于SSM的顺序约束而丢失的任何内容。然后，我们连接两个分支的输出，并通过一个最终的线性层进行投影。这种组合确保了最终的特征表示同时结合了顺序和空间信息，利用了两个分支的优势。我们注意到，每个分支的输出被投影到一个大小为$\frac{C}{2}$的嵌入空间中（即原始嵌入维度的一半），以保持与原始块设计相似的参数数量。给定输入$X_{\text{in}}$，MambaVision混合器$X_{\text{out}}$的输出计算如下：
$$\begin{array}{l}{X}_1=\mathrm{Scan}\left(\sigma \left(\mathrm{Conv}\left(\mathrm{Linear}\left(C,\frac{C}{2}\right)\left(X_{\text{in }}\right)\right)\right)\right)\\ X_2=\sigma \left(\mathrm{Conv}\left(\mathrm{Linear}\left(C,\frac{C}{2}\right)\left(X_{in}\right)\right)\right)\\ X_{\text{out }}=\mathrm{Linear}\left(\frac{C}{2},C\right)\left(\mathrm{Concat}\left(X_1,X_2\right)\right)\end{array}$$
其中，$\mathrm{Linear}\left(C_{\text{in}},C_{\text{out}}\right)(\cdot )$表示具有输入和输出嵌入维度$C_{\text{in}}$和$C_{\text{out}}$的线性层，$\mathrm{Scan}$是[2]中的选择性扫描操作，$\sigma$是激活函数，使用Sigmoid线性单元（SiLU）[17]。此外，$\mathrm{Conv}$和$\mathrm{Concat}$分别表示一维卷积和连接操作。在算法1中，我们提供了类似PyTorch的MambaVision混合器的伪代码。一般来说，我们提出的修改导致了更丰富的特征表示、更好的泛化能力以及计算机视觉任务上性能的提升。我们还在第X节中通过实验验证了每个设计选择的有效性。

自注意力 我们使用通用的多头自注意力机制，如下所示：
$$\mathrm{Attention}(Q,K,V)=\mathrm{Softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_h}}\right)V$$
其中，$Q$、$K$、$V$分别表示查询、键和值，$d_h$是注意力头的数量。不失一般性，注意力公式可以以类似于先前工作[10, 18]的窗口化方式计算。

4 实验

在ImageNet-1K数据集[19]上进行图像分类实验。我们遵循先前工作[10, 12, 29]的标准训练方案，以便对不同模型的性能进行可比分析。具体来说，所有模型均训练了300个epoch，使用余弦衰减调度器，该调度器分别为预热和冷却阶段额外使用20个epoch。此外，我们使用LAMB优化器[30]，全局批量大小为4096，初始学习率为0.005，权重衰减为0.05。我们注意到，与传统AdamW[31]相比，使用LAMB优化器取得了更好的结果，这尤其得益于其对较高学习率的鲁棒性。我们使用32块A100 GPU进行分类任务。

为了评估下游任务的性能，我们将预训练模型用作目标检测、实例分割和语义分割任务的骨干网络，并分别使用MS COCO数据集[32]和ADE20K数据集[33]。具体来说，对于目标检测和实例分割，我们使用Mask-RCNN[34]头，其超参数如$\times 3$学习率调度（初始学习率为0.0001）、批量大小为16和权重衰减为0.05。对于语义分割，我们使用UperNet网络[35]头，采用Adam-W[31]优化器，初始学习率为$6\times 10^{-5}$，全局批量大小为16。我们使用8块A100 GPU进行所有下游任务。

5 结果

5.1 图像分类

在表1中，我们展示了ImageNet-1K分类结果。具体来说，我们将模型与不同系列的模型（如基于卷积的、基于Transformer的、卷积-Transformer混合的和基于Mamba的）进行比较，并证明在考虑ImageNet Top-1准确率和图像吞吐量的情况下，我们的模型显著优于先前的工作。例如，与流行的模型如ConvNeXt和Swin Transformers相比，MambaVision-B（84.2%）优于ConvNeXt-B（83.8%）和SwinB（83.5%），同时还具有显著更好的图像吞吐量。与基于Mamba的模型相比，我们也观察到了类似的趋势。具体来说，MambaVision-B（84.2%）在图像吞吐量显著更高的情况下优于VMamba-B（83.9%）。我们还注意到，尽管我们的主要设计目标是准确率和吞吐量的权衡，但与类似规模的对应模型相比，MambaVision模型变体的FLOPs要低得多。例如，MambaVision-B的GFLOPs比MaxViT-B少56%。

5.2 目标检测和分割

我们在表2中展示了MS COCO数据集[32]上的目标检测和实例分割结果。具体来说，我们训练了不同检测规模的模型，以进一步验证MambaVision在不同场景下的有效性。我们还注意到，我们的目标不是在这些任务中实现最先进（SOTA）性能，而是将我们的骨干网络与规模相当的流行视觉骨干网络进行比较，并验证其有效性。使用简单的Mask-RCNN检测头，我们的预训练MambaVision-T骨干网络在框AP和掩码AP方面分别达到了46.4和41.8，优于ConvNeXt-T[20]和Swin-T[10]模型。使用Cascade Mask-RCNN网络，MambaVision-T、MambaVision-S和MambaVision-B均优于竞争对手。具体来说，在框AP和掩码AP方面，MambaVision模型分别优于ConvNeXt-T（+0.6和+0.6）、ConvNeXt-S（+0.2和+0.2）和ConvNeXt-B（+0.1和+0.1）。类似地，在框AP和掩码AP方面，MambaVision分别优于Swin-T（+0.6和+0.6）、Swin-S（+0.1和+0.2）和Swin-B（+0.9和+0.7）。

我们在表3中展示了ADE20K数据集[33]上的语义分割基准。对于这些实验，我们使用UPerNet[35]以与其他模型进行比较。我们观察到，对于不同变体，MambaVision模型均优于规模相当的竞争模型。例如，在mIoU方面，MambaVision-T、MambaVision-S和MambaVision-B分别优于Swin-T、Swin-S和Swin-B（+0.6、+0.6和+1.0）。尽管我们没有对下游任务的超参数调优进行大量优化，但这些结果证明了MambaVision作为不同视觉任务骨干网络的可行性，尤其是在高分辨率设置中。

5.3 消融实验

令牌混合器的设计：在本节中，我们对MambaVision令牌混合器进行系统设计的全面消融研究。我们的目标是修改现有的Mamba块以用于计算机视觉任务，并评估其在不同任务（如分类、目标检测、实例分割和语义分割）中的性能。所有实验均基于与MambaVision-T模型相似的架构布局。

如表4所示，我们首先使用原始Mamba公式，在SSM分支中使用因果卷积层（即conv1），而在我们提出的对称分支中不使用额外的卷积层（即conv2）。不出所料，这种公式实现了次优性能，Top-1准确率为80.9%（-1.8%），框AP和掩码AP分别为44.8（-1.6）和40.2（-1.6），mIoU为44.2%（-1.4%）。然后，我们将SSM分支的因果卷积（即conv1）替换为常规卷积层，并观察到由于这一变化，所有指标均有所改善。此外，我们添加了conv2层，但使用Mamba的相同门控机制而不是连接。这一变化提高了性能，Top-1准确率为81.3%，框AP和掩码AP分别为45.3和41.0，mIoU为45.7%。最后，添加连接显著提高了所有指标的性能，在ImageNet Top-1、MS COCO的框AP和掩码AP以及ADE20K数据集的mIoU方面分别提高了+1.0%、+1.1%、+0.8%和+0.9%。因此，这验证了我们的假设，即连接两个分支（即SSM和非SSM）的输出有助于学习更丰富的特征表示并增强全局上下文理解。

混合模式：在本节中，我们全面研究了自注意力和MambaVision令牌混合器的不同混合集成模式的影响。对于所有实验，架构遵循MambaVision-T布局，并且我们保持模型等参数以进行公平比较。这些模式用于具有混合功能的阶段3和阶段4。我们从使用随机模式开始，实现了次优的Top-1准确率81.3%。这证实了我们之前的直觉，即简单地使用自注意力而不采用特定模式可能并不有效。然后，我们在每个阶段的前$N/2$层中使用自注意力块（$N$表示阶段层的总数），并观察到Top-1准确率提高了+0.2%（81.5%）。然而，使用自注意力/MambaVision混合器的混合层模式会使准确率略微降低-0.1%（81.4%）。相反，通过使用MambaVision/自注意力反转混合层的顺序，性能得到了提高，并达到了Top-1准确率81.6%。然后，我们仅在每个阶段的最后$N/4$层中使用自注意力块，并观察到准确率立即提高了+0.3%（81.9%）。这验证了我们的假设，即在每个阶段的最后几层中使用自注意力块是一种有效的设计。然而，其表示学习能力需要根据MambaVision层进行调整。将自注意力块的数量增加到每个阶段的最后$N/2$层时，实现了最佳性能82.3%。

6 结论

在这项工作中，我们介绍了MambaVision，这是第一个专门为视觉应用设计的Mamba-Transformer混合骨干网络。我们提出了对Mamba公式的重新设计，以增强全局上下文表示学习能力，并呈现了混合设计集成模式的全面研究。MambaVision在Top-1准确率和图像吞吐量方面实现了新的SOTA帕累托前沿，显著优于基于Transformer和Mamba的模型。我们希望这些发现能够成为一类新的混合视觉模型的基础。