CVPR2025，英伟达提出首个Mamba-Transformer视觉骨干网络！（附论文和源码）
仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
以下文章来源于微信公众号：新智元
链接：https ://mp.weixin.qq.com/s/VFS0MVqSluzJ11h1zYuo-A
本文仅用于学术分享，如有侵权，请联系后台作删文处理
导读
论文链接：https ://arxiv.org/abs/2407.08083
代码链接：https ://github.com/NVlabs/MambaVision
正如标题所言「Attentionisallyouneed」，Transformer已成为不同领域的「霸主」，包括计算机视觉、自然语言处理、语音处理和机器人技术。
第一个挑战Transformer架构的是Mamba，一种新的状态空间模型（SSM），它具有线性时间复杂度，并在多个语言建模任务中超越或与Transformer媲美。
但在不同的视觉任务上，VisionTransformer(ViT)和卷积神经网络(CNN)架构的骨干网络，仍然优于基于Mamba的视觉模型。
而这一次，英伟达高级工程师AliHatamizade，宣布被顶会CVPR-2025接受的MambaVision，在视觉任务上超越以往的模型，而设计的关键在于将Mamba和Transformer混合。
正如图1所示，在ImageNet-1K基准上，MambaVision的Top-1准确率和图像处理能力达到了新的Pareto最优点，超越了Mamba、CNN和ViT基于的模型，有时差距非常显著。
在下游任务如目标检测、实例分割以及语义分割中，采用MambaVision作为骨干网络的模型在MSCOCO和ADE20数据集上分别超越了同等规模的对比模型。
MambaVision是首个针对计算机视觉应用，结合Mamba和Transformer的混合架构的尝试。主要贡献总结如下：
1引入了重新设计的适用于视觉任务的Mamba模块，提升了相较于原始Mamba架构的准确性和图像处理能力。
2系统性地研究了Mamba和Transformer模块的融合模式，并展示了在最终阶段加入自注意力模块，显著提高了模型捕捉全局上下文和长距离空间依赖的能力。
论文链接:https ://arxiv.org/abs/2407.08083
在这项工作中，作者系统地重新设计了Mamba模块，使其更加适合视觉任务。
新方法是一种混合架构，结合了新提出的公式（即MambaVisionMixer和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表示批归一化（BatchNormalization）；GELU是激活函数，表示GaussianErrorLinearUnit；z^是经过卷积、批归一化和激活函数处理后的中间结果；最后，z是通过卷积和批归一化后的结果与原始输入相加，形成残差连接。
这种结构有助于缓解深层网络训练中的梯度消失问题，并提高模型的训练效率。
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混合模块的选择。
在不失一般性的情况下，层归一化（LayerNormalization）被用于NormNorm。给定N层，前N/2层使用MambaVision混合模块，而剩余的N/2层使用自注意力机制。
MambaVision混合模块:重新设计了原始的Mamba混合模块，使其更适合视觉任务。
如图3所示，首先将因果卷积（causalconvolution）替换为常规卷积，因为因果卷积将信息限制在一个方向上，这对视觉任务来说不仅没必要，而且局限性还很大。
此外，添加了一个不包含SSM（状态空间模型）的对称分支，该分支由额外的卷积和SiLU激活函数组成，以补偿由于SSM的顺序约束而可能丢失的内容。
然后，将两个分支的输出拼接起来，并通过最终的线性层进行投影。这种组合确保了最终的特征表示，同时包含顺序信息和空间信息，从而充分利用了两个分支的优势。
注意到，每个分支的输出被投影到一个大小为C/2的嵌入空间（即原始嵌入维度的一半），以保持与原始模块设计相似的参数量。
给定输入Xin，MambaVision混合模块的输出Xout计算如下：
实验结果
表1展示了ImageNet-1K分类结果。具体来说，与不同类别的模型进行了比较，包括基于卷积的模型、基于Transformer的模型、卷积-Transformer混合模型以及基于Mamba的模型，并证明新模型在ImageNetTop-1准确率和图像处理能力方面大幅超越了之前的工作。
在与基于Mamba的模型比较时也观察到了类似的趋势。具体来说，尽管MambaVision-B（84.2%的图像处理能力显著更高，但仍优于VMamba-B（83.9%）。
与同等规模的模型相比，MambaVision型变体的FLOPs远低于它们。例如，MambaVision-B的GFLOPs比MaxViT-B少了56%。
表2展示在MSCOCO数据集上的目标检测和实例分割结果。
具体来说，训练了不同检测尺寸的模型，以进一步验证MambaVision不同场景下的有效性。
通过简单的Mask-RCNN检测头,预训练的MambaVision-T骨干网络，超过了ConvNeXt-T和Swin-T模型。
使用CascadeMask-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 ://x.com/ahatamiz1/status/1894838302450295183
欢迎加入《AI未来星球》，一起成长
扫描下方二维码即可加入~
你可以获得什么？
1、大白之前花费10W+购买，AI行业各场景私有数据集下载，星球内倾情分享；2、AI行业研发、产品、商业落地问题咨询（目前AI公司创业中），都可获高质量解答，有效期一年，无限次提问，有问必答。3、定期邀请AI行业各类嘉宾分享，创业/商业等方面的经验！
帮助你解决遇到的实际问题，升职加薪！
大家一起加油！
