英伟达提出首个Mamba-Transformer视觉骨干网络！打破精度/吞吐瓶颈 | CVPR 2025

新智元报道

编辑：KingHZ

【新智元导读】CVPR 2025，混合新架构MambaVision来了！Mamba+Transformer混合架构专门为CV应用设计。MambaVision 在Top-1精度和图像吞吐量方面实现了新的SOTA，显著超越了基于Transformer和Mamba的模型。

正如标题所言「Attention is all you need」，Transformer已成为不同领域的「霸主」，包括计算机视觉、自然语言处理、语音处理和机器人技术。

第一个挑战Transformer架构的是Mamba，一种新的状态空间模型（SSM），它具有线性时间复杂度，并在多个语言建模任务中超越或与Transformer媲美。

但在不同的视觉任务上，Vision Transformer (ViT) 和卷积神经网络 (CNN) 架构的骨干网络，仍然优于基于Mamba的视觉模型。

而这一次，英伟达高级工程师Ali Hatamizade，宣布被顶会CVPR-2025接受的MambaVision，在视觉任务上超越以往的模型，而设计的关键在于将Mamba和Transformer混合。

正如图1所示，在ImageNet-1K基准上，MambaVision的Top-1准确率和图像处理能力达到了新的Pareto最优点，超越了Mamba、CNN和ViT基于的模型，有时差距非常显著。

在下游任务如目标检测、实例分割以及语义分割中，采用MambaVision作为骨干网络的模型在MS COCO和ADE20数据集上分别超越了同等规模的对比模型。

MambaVision是首个针对计算机视觉应用，结合Mamba和Transformer的混合架构的尝试。主要贡献总结如下：

1 引入了重新设计的适用于视觉任务的Mamba模块，提升了相较于原始Mamba架构的准确性和图像处理能力。

2 系统性地研究了Mamba和Transformer模块的融合模式，并展示了在最终阶段加入自注意力模块，显著提高了模型捕捉全局上下文和长距离空间依赖的能力。

论文链接:https://arxiv.org/abs/2407.08083

在这项工作中，作者系统地重新设计了Mamba模块，使其更加适合视觉任务。

新方法是一种混合架构，结合了新提出的公式（即MambaVision Mixer和MLP）以及Transformer模块。

具体来说，研究了不同的集成模式，比如以等参数方式将Transformer模块添加到早期、中间和最终层，或者每隔l层添加一次。

分析表明，在最终阶段利用多个自注意力模块，可以显著增强捕捉全局上下文和长程空间依赖的能力。

使用混合架构相较于纯Mamba或ViT模型，图像处理能力也得到了显著提升。

网络架构

宏观架构

如图2所示，MambaVision 采用了分层架构，由4个不同的阶段组成。

前两个阶段使用基于CNN的层，负责在较高输入分辨率下进行快速特征提取，而第3和第4阶段则包括了新提出的 MambaVision和Transformer模块。

具体来说，给定一个大小为H×W×3的图像，输入首先被转换为大小为H/4×W/4×C的重叠patch，并通过两层连续 3×3的CNN 层（步幅为2）构成的主干投影到C维嵌入空间中。

在各个阶段之间的下采样模块由一个批归一化的3×3的CNN 层（步幅为2）组成，将图像分辨率减半。

此外，第1和第2阶段中的CNN模块，采用了通用的残差模块结构，具体如下：

其中：Conv3×3 表示3×3卷积操作；BN表示批归一化（Batch Normalization）；GELU 是激活函数，表示 Gaussian Error Linear Unit；z^ 是经过卷积、批归一化和激活函数处理后的中间结果；最后，z是通过卷积和批归一化后的结果与原始输入相加，形成残差连接。

这种结构有助于缓解深层网络训练中的梯度消失问题，并提高模型的训练效率。

Mamba架构

Mamba是结构化状态空间序列模型的扩展，能够通过可学习的隐状态 h(t)，将一维连续输入x(t)转换为y(t)。该过程的公式如下：

其中，矩阵A，B，C是模型的参数。

离散化：为了提高计算效率，以上公式中的连续参数A，B和C需要转化为离散参数。具体而言，假设时间步长为Δ，可以应用零阶保持规则来获取离散参数：

这种离散化方法能够提升计算效率，便于在实际应用中实现 Mamba 模型。

使用离散参数代入到原方程：

此外，对于一个大小为T的输入序列，可以用带有卷积核K的全局卷积，进一步简化上式中的输出，具体如下

选择性：Mamba 进一步扩展了S4公式，引入了一种选择机制，使得模型能够进行依赖于输入的序列处理。这种机制使得模型的参数B 、C和Δ可以根据输入动态调整，从而滤除无关信息。

设输入X是TxC矩阵，其中 T 为序列长度，C为嵌入维度，第3和第4阶段的第n层输出可以按如下方式计算：

其中，NormNorm和MixerMixer分别表示层归一化和 token 混合模块的选择。

层架构

在不失一般性的情况下，层归一化（Layer Normalization）被用于 NormNorm。给定N层，前 N/2层使用 MambaVision混合模块，而剩余的N/2层使用自注意力机制。

MambaVision 混合模块:重新设计了原始的Mamba混合模块，使其更适合视觉任务。

如图3所示， 首先将因果卷积（causal convolution）替换为常规卷积，因为因果卷积将信息限制在一个方向上，这对视觉任务来说不仅没必要，而且局限性还很大。

此外，添加了一个不包含SSM（状态空间模型）的对称分支，该分支由额外的卷积和SiLU激活函数组成，以补偿由于SSM的顺序约束而可能丢失的内容。

然后，将两个分支的输出拼接起来，并通过最终的线性层进行投影。这种组合确保了最终的特征表示，同时包含顺序信息和空间信息，从而充分利用了两个分支的优势。

注意到，每个分支的输出被投影到一个大小为C/2的嵌入空间（即原始嵌入维度的一半），以保持与原始模块设计相似的参数量。

给定输入Xin，MambaVision混合模块的输出Xout计算如下：

其中，Linear(Cin,Cout)(⋅)表示一个线性层，输入和输出的嵌入维度分别为Cin和Cout；Scan是选择性扫描操作（selective scan）；σ是激活函数，这里使用的是Sigmoid线性单元（SiLU；Conv和Concat 分别表示1D卷积和拼接操作。

实验结果

表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%。

表2展示在MS COCO数据集上的目标检测和实例分割结果。

具体来说，训练了不同检测尺寸的模型，以进一步验证 MambaVision 不同场景下的有效性。

通过简单的Mask-RCNN检测头,预训练的MambaVision-T骨干网络，超过了 ConvNeXt-T和 Swin-T模型。

使用Cascade Mask-RCNN网络时，MambaVision-T、MambaVision-S和MambaVision-B都超过了竞争对手。

表3展示了在ADE20K数据集上的语义分割基准测试。

对于这些实验，使用了 UPerNet，以便与其他模型进行比较。

观察到，MambaVision 模型在不同变体下超越了同等规模的竞争模型。

例如，MambaVision-T、MambaVision-S 和 MambaVision-B分别在mIoU上超越了Swin-T、Swin-S和Swin-B，提升幅度为+0.6、+0.6和+1.0。

尽管没有对下游任务进行大量的超参数调优优化，这些结果仍然证明了MambaVision作为一种有前景的视觉任务骨干网络的可行性，特别是在高分辨率设置下。

消融实验和更多细节请参考原文。

参考资料：

https://arxiv.org/abs/2407.08083

https://x.com/ahatamiz1/status/1894838302450295183