【新智元导读】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/statu...
END
作者：新智元
来源：AIWalker
推荐阅读
- 自回归视觉生成破局：Next-X 预测，开辟视觉新路径
- YOLOv12 来袭，更高性能、更快速度~
- 从文字到视觉：EvalMuse-40K 如何评价 T2I 模型的进化
- FoundIR: 释放百万规模高质量训练数据，助推图像复原基础模型
- CPU推理1ms的SHViT Backbone 来啦
本文章著作权归作者所有，任何形式的转载都请注明出处。更多动态滤波，图像质量，超分辨相关请关注我的专栏AIWalker。欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
