英伟达也对 Mamba下手了 ，视觉 Transformer 与 Mamba 的完美融合 ！

作者提出了一种新颖的混合Mamba-Transformer架构，称为MambaVision，这是专门为视觉应用量身定制的。作者的核心贡献包括重新设计Mamba公式，以增强其高效建模视觉特征的能力。此外，作者还对将视觉Transformer（ViT）与Mamba集成的可行性进行了全面的消融研究。作者的结果表明，在Mamba架构的最后几层配备几个自注意力块，大大提高了捕获长距离空间依赖关系的建模能力。基于作者的发现，作者引入了一系列具有分层架构的MambaVision模型，以满足各种设计标准。在ImageNet-1K数据集上的图像分类中，MambaVision模型变体在Top-1准确率和图像吞吐量方面达到了新的最先进（SOTA）性能。在MS COCO和ADE20K数据集上的下游任务，如目标检测、实例分割和语义分割中，MambaVision超越了同等大小的架构，并展示了更优的性能。

代码：https://github.com/NVIabs/MambaVision。

1 Introduction

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

图1：ImageNet-1K数据集上的Top-1准确性与图像吞吐量。 所有测量均在A100 GPU上进行，批量大小为128。MambaVision达到了新的SOTA帕累托前沿。

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

Vision Mamba (Vim) [3] 等人提出了修改措施，例如双向SSM来解决全局上下文和空间理解的不足。尽管双向SSM有潜力捕获更全面的上下文，但由于需要在做出预测之前处理整个序列，它们引入了显著的延迟。此外，复杂度的增加可能导致训练困难、过拟合的风险，并且不一定能带来更高的准确度。由于这些缺陷，带有Vision Transformer (ViT) 和卷积神经网络 (CNN) 架构的 Backbone 网在不同的视觉任务上仍然优于最佳的基于Mamba的视觉模型。

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

作者引入了MambaVision模型，该模型由多分辨率架构组成，并利用基于CNN的残差块快速提取较大分辨率特征。如图1所示，MambaVision在ImageNet-1K Top-1准确性和图像吞吐量方面达到了新的SOTA帕累托前沿，超过了Mamba、CNN和基于ViT的模型，有时幅度还很大。在如下游任务目标检测和实例分割以及语义分割中，带有MambaVision Backbone 网的模型在MS COCO和ADE20数据集上分别优于同等大小的对应模型。因此，这验证了MambaVision作为有效 Backbone 网的有效性和多功能性。

据作者所知，MambaVision是首次研究和开发同时包含Mamba和Transformers的混合架构以用于计算机视觉应用。

作者在这项工作中的主要贡献总结如下：

• 作者引入了一个重新设计的面向视觉的Mamba块，提高了原始Mamba架构的准确性和图像吞吐量。
• 作者系统地调查了Mamba和Transformer块的集成模式，并证明在最后阶段整合自注意力块显著提高了模型捕获全局上下文和长距离空间依赖的能力。
• 作者介绍了MambaVision，这是一个新颖的混合Mamba Transformer模型。分层的MambaVision在ImageNet-1K数据集上实现了Top-1和图像吞吐量折衷的新SOTA帕累托前沿。

2 Related work

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

Mamba： 自从Mamba被引入以来，已经提出了许多努力来利用其在视觉应用中的能力。具体来说，Vim[3]提出使用双向SSM公式，与相同的Mamba公式相同，在这种公式中，标记在前后两个方向上处理，以捕捉更多的全局上下文并提高空间理解。然而，双向编码增加了计算负载，可能会减慢训练和推理时间。此外，有效结合来自多个方向的信息形成一个连贯的全局理解是具有挑战性的，因为在过程中可能会丢失一些全局上下文。与Vim相比，作者提出的MambaVision使用单一的前向传播和重新设计的Mamba块，可以捕捉到短距离和长距离的信息，并且在ImageNet Top-1准确性和吞吐量方面显著优于Vim。

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

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

3 Methodology

Macro Architecture

在本节中, 作者介绍了MambaVision, 这是作者提出的新架构, 在ImageNet-1K数据集上具有 SOTA性能。如图2所示, MambaVision具有分层架构, 由4个不同阶段组成。前两个阶段由基于CNN的层组成, 用于在更高的输入分辨率下快速特征提取, 而第3和4阶段包括作者提出的 MambaVision和Transformer块。具体来说, 给定一个大小为 的图像, 输入首先被转换成大小为 的重叠图像块, 并通过干细胞（由两个连续的步长为 2 的 CNN 层组成）投射到 维嵌入空间中。阶段之间的下采样器由一个批标准化步长为 2 的 CNN 层组成, 它将图像分辨率减少一半。此外, 第 1 和第 2 阶段的CNN块遵循以下通用残差块公式

注意：表格内容在翻译中已调整为适应中文语境的顺序。公式和引用编号保持不变。

Micro Architecture

在本节中，作者首先回顾了Mamba和SSMs的基础知识。然后作者分阶段展示了第3和第4阶段的微架构设计，并详细讨论了MambaVision的公式化表述。

3.2.1 Mamba Preliminaries

Mamba是结构化状态空间序列模型（S4）的扩展, 能够通过可学习的隐藏状态 将 1D连续输入 转换为 , 其参数为 和 , 根据以下公式：

离散化上述公式中的连续参数 和 C被进一步转换为离散参数以提高计算效率[16]。具体而言, 假设一个时间尺度 , 可以应用零阶保持规则来获得离散参数 和 , 根据以下公式:

然后方程式2可以用离散参数表示为：

此外，对于大小为 的输入序列，可以应用全局卷积核页来计算方程式4的输出，如下所示：

图3：MambaVision块的架构。除了用它们的常规对应物替换因果卷积层外，作者还创建了一个没有SSM的对称路径作为标记混合器，以增强全局上下文的建模。

图2：分层MambaVision模型的架构。前两个阶段使用残差卷积块进行快速特征提取。第3和第4阶段同时使用MambaVision和Transformer块。具体来说, 给定 层, 作者使用 MambaVision和MLP块, 后面跟着额外的 Transformer和MLP块。最终层中的Transformer块允许恢复丢失的全局上下文并捕捉长距离空间依赖。

算法1 类似PyTorch的伪代码用于MambaVision混合器

选择性Mamba通过引入选择机制进一步扩展了S4公式，该机制允许进行输入相关的序列处理。这使得模型的参数，和可以根据输入动态调整，并过滤掉不相关的信息。进一步的离散化细节在[2]中提供。

3.2.2 Layer Architecture

假设输入 , 其中序列长度为 , 嵌入维度为 , 第 3 和第 4 阶段的第 层输出可以按如下方式计算:

其中 Norm 和 Mixer 分别表示层归一化和标记混合块的选择。不失一般性, Norm 使用层归一化。给定 层, 前 层采用 MambaVision 混合块, 而剩余的 层采用自注意力。下面作者将描述每个混合块的具体细节。

MambaVision 混合块如图3所示，作者重新设计了原始的 Mamba 混合块，使其更适合视觉任务。首先，作者建议用常规卷积替换因果卷积，因为它限制了单向的影响，这对于视觉任务是不必要且受限的。此外，作者增加了一个没有 SSM 的对称分支，该分支包含一个额外的卷积和 SiLU 激活，以补偿由于 SSMs 的序列限制而造成的任何内容丢失。然后作者将两个分支的输出连接起来并通过一个最终的线性层进行投射。这种组合确保了最终的特征表示融合了序列和空间信息, 利用了两个分支的优势。作者注意到, 每个分支的输出被投射到一个大小为 (即原始嵌入维度的一半) 的嵌入空间, 以保持与原始块设计相似的参数数量。给定输入 , MambaVision 混合块的输出 计算如下:

其中 Linear 表示一个具有 和 作为输入和输出嵌入维度的线性层, Scan 是如 [2] 中所示的选择性扫描操作, 是激活函数, 使用的是 Sigmoid Linear Unit (SiLU) [17]。此外, Conv 和 Concat 分别表示1D卷积和连接操作。在算法1中, 作者提供了一个类似 PyTorch 的伪代码用于 MambaVision 混合块。总的来说, 作者提出的修改导致特征表示更丰富, 泛化能力更强, 并在计算机视觉任务上提高了性能。作者也实验验证了每个设计选择的有效性, 这部分内容在 Sec.

自注意力作者使用一种通用的多头自注意力机制，按照以下方式计算：

分别表示 Query 、键和值, 是注意力头的数量。不失一般性, 可以以窗口方式计算注意力公式，类似于之前的工作 。

4 Experiments

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

为了评估下游任务的性能，作者将预训练的模型作为目标检测和实例分割以及语义分割任务的 Backbone 网络，并分别使用了MS COCO数据集[32]和ADE20K数据集[33]。具体来说，对于目标检测和实例分割，作者使用了Mask-RCNN[34] Head ，并设置了如 LR计划、初始学习率为0.0001、批量大小为16、权重衰减为0.05等超参数。对于语义分割，作者使用了UperNet网络[35] Head 和Adam-W[31]优化器，初始学习率为6e-5，全局批量大小为16。作者在所有下游任务中使用了8个A100 GPU。

5 Results

Image classification

在表1中，作者展示了ImageNet-1K分类的结果。具体来说，作者与不同类型的模型家族进行了比较，如基于卷积的、基于Transformer的、卷积-Transformer混合的以及基于Mamba的模型，并展示了作者的模型在ImageNet Top-1准确率和图像吞吐量方面超过了此前的成果。例如，与流行的模型如ConvNeXt和Swin Transformers相比，MambaVision-B（84.2%）分别超过了ConvNeXt-B（83.8%）和Swin-B（83.5%），同时在图像吞吐量上也有显著优势。与基于Mamba的模型相比，作者也观察到了类似的趋势。具体来说，MambaVision-B（84.2%）尽管图像吞吐量明显更高，但性能还是超过了V Mamba-B（83.9%）。作者还想指出，虽然作者主要的设计目标是准确性与吞吐量的权衡，但MambaVision模型变体与同等大小的模型相比，FLOPs要低得多。例如，MambaVision-B的GFLOPs比MaxViT-B少了。

Object Detection and Segmentation

作者在表2中展示了在MS COCO数据集[32]上的目标检测和实例分割结果。特别地，作者训练了不同检测尺寸的模型，以进一步验证MambaVision在不同场景下的有效性。作者还注意到，作者的目标并非要在这些任务上达到最先进水平，而是要比较作者的主干网络与同尺寸的流行视觉主干网络的表现，并验证其有效性。使用简单的Mask-RCNN检测Head，作者预训练的MambaVision-T主干网络在box AP和mask AP方面分别达到了46.4和41.8，超过了ConvNeXt-T [20]和Swin-T [10]模型。使用级联Mask-RCNN网络，MambaVision-T、MambaVision-S和MambaVision-B均超过了竞争模型。具体来说，在box AP和mask AP方面，MambaVision模型分别比ConvNeXt-T高出+0.2和0.2，比ConvNeXt-B高出+0.1和0.1。同样，在box AP和mask 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作为一个不同视觉任务的有前景的主干网络是可行的，尤其是在高分辨率设置下。

Ablation

在本节中, 作者对MambaVision标记混合器进行了全面的消融研究, 以系统化设计。作者的目标是修改现有的Mamba模块以用于计算机视觉任务, 并评估其在分类、目标检测、实例分割和语义分割等不同任务上的性能。所有实验都遵循类似于MambaVision-T模型架构布局的模型。如表4所示, 作者首先使用SSM分支 (即conv1) 和对称分支（即conv2）中的因果卷积层的原始Mamba公式。正如预期的那样, 这种公式取得了次优的性能, Top-1准确率为 ( , box AP和mask AP分别为 和 为 ）。然后, 作者将SSM分支中的因果卷积（即conv1）替换为常规卷积层, 并观察到由于这种改变, 所有指标都有所改进。此外, 作者增加了conv2层, 但使用了Mamba的相同门控机制,而不是拼接。这一改变提高了性能, 使得Top-1准确率达到 , box AP和mask AP分别为 45.3 和 为 。最后, 通过拼接显著提高了所有指标的性能, 分别提高了 和 +0.9 , 这些指标分别是ImageNet Top-1、MS COCO的box AP和mask AP以及ADE2OK数据集的mloU。因此, 这验证了作者的假设, 即拼接两个分支（即SSM和非 SSM）的输出可以学习更丰富的特征表示并增强全局上下文理解。

混合模式在本节中, 作者全面研究了自关注和MambaVision标记混合器的不同混合整合模式的影响。对于所有实验, 架构遵循MambaVision-T布局, 并保持模型等参数以公平比较。这些模式在第3和第4阶段使用混合功能。作者首先使用随机模式, 取得了次优的Top-1准确率 。这证实了作者之前的直觉, 即简单地使用自关注而没有特定的模式可能并不有效。

然后, 作者在每个阶段的前 层使用自关注块, 其中 表示阶段层的总数, 并观察到Top-1 准确率提高了 。

然而, 使用自关注/MambaVision混合器块的混合层模式略微降低了准确率, 为 ( ）。相反, 通过使用MambaVision/自关注混合层的反转顺序提高了性能, 达到了Top-1 准确率 ）。然后, 作者只在每个阶段的最后 层使用自关注块, 并立即观察到准确率提高了 。这验证了作者的假设, 即在每个阶段的最后几层使用自关注块是一种有效的设汁。然而, 其表示学习能力需要相对于MambaVision层进行调整。将自关注块的数量增加到每个阶段的最后 层, 取得了最佳的 性能。

6 Conclusion

在本文中，作者引入了MambaVision，这是首个专门为视觉应用设计的Mamba-Transformer混合骨架。

作者提出了重新设计Mamba公式的方法，以增强全局上下文表示的学习能力，并进行了混合设计集成模式的综合研究。

MambaVision在Top-1准确性和图像吞吐量方面达到了新的SOTA帕累托前沿，大幅超越了基于Transformer和Mamba的模型。

作者希望这些研究成果能够成为新型混合视觉模型的基础。

参考

[1].MambaVision: A Hybrid Mamba-Transformer Vision Backbone.

公众号后台回复“数据集”获取100+深度学习各方向资源整理