视觉任务中真的需要 Mamba吗 ？

最近引入的Mamba架构，具有状态空间模型（SSM）类似RNN的标记混合器，旨在解决注意力机制的二次复杂度问题，并随后应用于视觉任务1。 然而，与卷积和基于注意力的模型相比，Mamba在视觉任务中的表现通常令人失望。在本文中，作者深入探讨了Mamba的本质，从概念上得出Mamba非常适合具有长序列和自回归特性的任务。 对于视觉任务，如图像分类并不符合这些特性，作者假设这个任务并不需要Mamba； 检测和分割任务也不具有自回归特性，但它们符合长序列特性，因此作者认为仍有必要探索这些任务中Mamba的潜力。 为了实证验证作者的假设，作者构建了一系列名为 MambaOut 的模型，通过堆叠 Mamba 块同时移除它们的核心标记混合器，SSM。实验结果强烈支持作者的假设。 具体来说，作者的MambaOut模型在ImageNet图像分类上超过了所有视觉Mamba模型，表明这个任务确实不需要Mamba。 对于检测和分割，MambaOut无法匹敌最先进的视觉Mamba模型的性能，这显示了Mamba在长序列视觉任务中的潜力。

1 引言

近年来，Transformer [76]已成为各种任务的主流 Backbone 网络，支撑着诸如BERT ，GPT系列 和ViT [23]等众多杰出模型。然而，Transformer的标记混合器——注意力[3]，随着序列长度的增加，其复杂度呈现二次增长，给长序列带来了重大挑战。

为了解决这个问题，引入了多种与标记长度成线性复杂度的标记混合器，例如动态卷积，Linformer ，Longformer ，Big Bird ，和Performer 。近来，一类类似RNN的新模型涌现出来，因其能在长序列上进行并行训练和高效推理而受到社区的广泛关注。

特别是，像RWKV 和Mamba这样的模型被证明在大语言模型（LLMs）中作为 Backbone 网络是有效的。

受到类似RNN模型的强大能力的启发，各种研究尝试将Mamba 引入到视觉识别任务中，例如Vision Mamba ，VMamba，LocalMamba ，和PlainMamba [88]等开创性工作。

Mamba的标记混合器是有结构的状体空间模型（SSM），遵循RNN的精神。

然而，他们的实验表明，基于SSM的视觉模型实际上与最先进的卷积模型和基于注意力的模型相比，性能令人失望。

这引出了一个引人注目的研究问题：作者在视觉任务中真的需要Mamba吗？

在本文中，作者探讨了Mamba的本质，并从概念上总结出Mamba理想适用于具有两个关键特征的任务：长序列和自回归，这是由于SSM的固有RNN机制[27, 26, 25]（见图2和图3的解释）。不幸的是，并不是很多视觉任务同时具备这两个特征。例如，在ImageNet上的图像分类符合这两个特征中的任何一个；

而在COCO上的目标检测和实例分割以及在ADE20K上的语义分割仅符合长序列特征。另一方面，自回归特征要求每个标记仅从前一个标记和当前标记聚合信息，这种概念在标记混合中被称为因果模式[63]（见图3(a)）。

实际上，所有视觉识别任务都属于理解领域，而不是生成领域，这意味着模型可以一次看到整个图像。因此，在视觉识别模型中对标记混合施加额外的因果约束可能导致性能下降（见图3(b)）。尽管这个问题可以通过双向分支[68]来缓解，但不可避免的是，这个问题在每个分支中仍然存在。

基于上述概念性讨论，作者提出了以下两个假设：

• 假设1： SSM对于图像分类并不是必要的，因为这项任务既不符合长序列也不符合自回归特征。
• 假设2： SSM可能对目标检测和实例分割以及语义分割有益，因为这些任务遵循长序列特征，尽管它们不是自回归的。

为了实验性地验证作者的假设，作者通过堆叠门控CNN[18]块开发了一系列名为_MambaOut_的模型。门控CNN和Mamba块之间的关键区别在于SSM的存在，如图1(a)所示。实验结果表明，更简单的MambaOut模型实际上已经超过了视觉Mamba模型[104, 50, 37, 88]的性能，这验证了作者的_假设1_。

作者还展示了实证结果，MambaOut在检测和分割任务上未能匹敌最先进的视觉Mamba模型[50, 37]的性能（见表2和表3），这强调了SSM在这些任务上的潜力，并有效地验证了作者的假设2。

作者的论文贡献有三方面。

1. 作者分析了SSM的类似RNN的机制，并从概念上得出Mamba适合处理具有长序列和自回归特性的任务。
2. 作者探讨了视觉任务的特点，并假设在ImageNet图像分类任务中不需要SSM，因为该任务不具备这两种特性，然而探索SSM在检测和分割任务中的潜力仍然有价值，因为这些任务符合长序列特性，尽管它们不是自回归的。
3. 作者开发了一系列名为MambaOut的模型，基于门控CNN块，但去掉了SSM。实验表明，在ImageNet图像分类任务中，MambaOut有效地超越了视觉Mamba模型，但在检测和分割任务中并未达到最先进的视

Mamba模型的性能。这些观察结果反过来验证了作者的假设。因此，由于其“奥卡姆剃刀”本质，MambaOut可能很容易作为未来视觉Mamba模型研究的一个自然 Baseline 。

2 Related work

Transformer已被广泛应用于各个领域，如NLP领域的BERT [20]和GPT系列[60; 61; 6; 1]，以及计算机视觉领域的ViT [23]。然而，Transformer中的注意力模块与序列长度成二次方关系，这带来了显著的计算挑战。

众多研究[72]探索了各种策略来减轻这个问题，包括低秩方法，核化， Token 混合范围限制，以及历史记忆压缩[62]。

近来，类似RNN的方法，特别是RWKV和Mamba，因其在大语言模型[59; 47]中取得的令人瞩目的成果而受到关注。

热情的探索性研究行人迅速将SSM和Mamba[25]融入到视觉识别任务中[104; 50; 37; 88; 44; 57; 58; 99; 86]。例如，Vision Mamba[104]集成了Mamba[25]来发展类似于ViT[23]的各向同性视觉模型；VMamba[50]使用Mamba构建类似于AlexNet[42]和ResNet[32]的分层视觉模型；

LocalMamba[37]通过融入局部诱导偏差来增强视觉Mamba模型[104; 50]；PlainMamba[88]旨在进一步提高各向同性Mamba模型的性能；

EfficientVMamba[58]通过引入孔洞选择性扫描来关注轻量级视觉Mamba模型的效率。

与这些举措不同，作者的工作并非旨在设计新的视觉Mamba模型。相反，作者探讨了一个关于在传统视觉识别情境中是否需要Mamba[25]的相关研究问题。

作者希望这篇论文能为未来对视觉Mamba模型的研究提供洞见。

3 Conceptual discussion

在本节中，作者首先讨论Mamba模型适合于哪些任务特征。接下来，作者检查视觉识别任务是否符合这些特征。根据检查结果，作者提出了关于视觉任务中对Mamba必要性的假设。

What tasks is Mamba suitable for?

Mamba的标记混合器是选择性的SSM [26; 25]，它定义了四个输入相关参数 并将它们通过以下方式转换为 ：

然后，SSM的序列到序列变换可以表示为：

其中 表示时间步， 代表输入， 表示隐藏状态， 指示输出。方程式2的递归性质[34]将类似RNN的SSM与因果注意力区分开来。隐藏状态 可以被视为一个固定大小的内存，存储所有历史信息。通过方程式2，这个记忆在更新时其大小保持不变。固定大小意味着记忆不可避免地是有损的，但它确保了将记忆与当前输入集成的计算复杂性保持恒定。相反，因果注意力将之前所有标记的键和值作为其记忆存储，通过添加每个新输入的标记的键和值来扩展。这种记忆在理论上是无损的。然而，随着更多标记的输入，记忆大小增长，从而增加了将记忆与当前输入集成的复杂性。RNN-like模型与因果注意力在记忆机制上的差异在图2中进一步说明。

由于SSM的记忆本质上是有损的，因此在处理短序列方面，Mamba无法展示其优势，而注意力在这种场景下可以轻松地表现良好。然而，在涉及长序列的场景中，由于注意力的二次复杂度，它的表现会下降。在这种情况下，Mamba可以明显突出其在将记忆与当前输入融合方面的效率，从而顺畅地处理长序列。因此，Mamba特别适合处理长序列。

尽管SSM的递归性质（方程式2）使Mamba能够有效地处理长序列，但它引入了一个重要的限制： 只能访问前一个和当前时间步的信息。如图3所示，这种标记混合被称为因果模式，可以表述为：

其中 和 分别表示第 个标记的输入和输出。由于其因果性质，这种模式非常适合自回归生成任务。

另一种模式称为完全可见模式，在这种模式下，每个标记可以从前面的所有标记和后续标记中聚合信息。这意味着每个标记的输出依赖于所有标记的输入：

其中 表示标记的总数。完全可见模式适合理解任务，在这些任务中，模型可以一次访问所有输入。

默认情况下，注意力处于完全可见模式，但通过将因果 Mask 应用于注意力图，它可以很容易地转变为因果模式。由于它们的递归性质，RNN-like模型本质上在因果模式下运行，如Mamba的方程式2所示。由于这种固有特性，RNN-like模型不能转换为完全可见模式。尽管RNN可以通过双向分支近似完全可见模式，但每个分支仍然单独保持在因果模式。因此，由于递归性质的固有限制，Mamba非常适合需要因果标记混合的任务。

总之，Mamba非常适合具有以下特点的任务：

• _特点1_：任务涉及处理长序列。
• _特点2_：任务需要因果标记混合模式。

接下来，作者将讨论视觉识别任务是否表现出这两个特点。

Do visual recognition tasks have very long sequences?

在本小节中，作者探讨了视觉识别任务是否需要长序列建模。作者以Transformer模型[76]为例来进行分析。考虑一个具有常见MLP比例4的Transformer块；假设其输入具有 Token 长度和通道（嵌入）维度，该块的FLOPs可以计算为：

从中，作者推导出中二次项与一次项的比例为：

如果，则中二次项的计算负载超过了线性项。这提供了一个简单的度量标准，以确定任务是否涉及长序列。例如，在ViT-S中具有384个通道时，阈值，在ViT-B中具有768个通道时，。

对于ImageNet上的图像分类，典型的输入图像大小为，在 Patch 大小为的情况下，产生个 Token 。显然，远小于和，这表明在ImageNet上的图像分类并不算作长序列任务。

对于COCO上的目标检测和实例分割，推理图像大小为，对于ADE20K上的语义分割，推理图像大小为，在 Patch 大小为的情况下， Token 数量约为4K。由于且，COCO上的检测和ADE20K上的分割都可以被认为是长序列任务。### 视觉识别任务需要因果 Token 混合模式吗？

如第3.1节所述，并在图3中说明，完全可见的 Token 混合模式允许不受限制的混合范围，而因果模式限制当前 Token 只从前面的 Token 访问信息。视觉识别被归类为理解任务，其中模型可以一次看到整个图像，从而消除了对 Token 混合的限制。对 Token 混合施加额外的限制可能会潜在地降低模型性能。如图3(b)所示，当将因果限制应用于Vision Transformers (ViT) [23; 73]时，观察到性能明显下降。通常，完全可见模式适用于理解任务，而因果模式更适合自回归任务。这一说法也可以通过观察BERT [20]和ViT [23]（BEiT [4]和MAE [30]）更多地用于理解任务，而不是GPT-1/2 [60; 61]和图像GPT [9]来证实。因此，视觉识别任务不需要因果 Token 混合模式。

Hypotheses regarding the necessity of Mamba for vision

根据作者之前的讨论，作者将关于在视觉识别任务中引入Mamba的必要性的假设总结如下：

• _假设1_：对于ImageNet上的图像分类任务，没有必要引入SSM，因为该任务不符合_特性1_或_特性2_。
• _假设2_：尽管不符合_特性1_，但由于视觉检测和分割任务符合_特性2_，进一步探索SSM的潜力仍然是值得的。

4 Experimental verification

Gated CNN and MambaOut

接下来，作者旨在实证验证作者的假设。

如图1(a)所示，Mamba块是基于Gated CNN块[18]的[25]。Gated CNN和Mamba的元架构都可以被视为MetaFormer[91]的 Token 混合器和MLP的简化集成，类似于MetaNeXt[93]。形式上，给定输入，元架构可以表述为：

其中表示归一化[38; 2; 83]；指的是进行 Token 混合的模块[92]；，和是可学习的权重。

比较方程10和11，并参考图1(a)，Gated CNN[60]和Mamba块[25]之间的主要区别在于SSM的存在。这促使作者开发了一系列基于Gated CNN块（无SSM）的模型，称为MambaOut。MambaOut将帮助作者评估Mamba对于视觉识别任务的必要性。

具体来说，作者将Gated CNN的 Token 混合器指定为核大小的深度卷积[11]，遵循ConvNeXt[52; 55]。此外，为了提高实际速度，作者只对部分通道进行深度卷积[54; 93; 7]，遵循InceptionNeXt[93]。如算法1所示，Gated CNN块的实现简单而优雅。与ResNet类似，作者采用4阶段框架，在每个阶段堆叠Gated CNN块来构建MambaOut，如图4所示。

Image classification on ImageNet

设置. ImageNet [19; 66] 作为图像分类的黄金标准基准测试，涵盖了1000个常见类别。它包括大约130万张训练图像和5万张验证图像。训练方案遵循DeiT [73] 不使用蒸馏。具体来说，所使用的数据增强包括随机调整大小的裁剪（输入图像大小为 ），水平翻转，RandAugment [15]，Mixup [100]，CutMix [96]，随机擦除 [102] 和颜色抖动，以及正则化技术包括权重衰减，随机深度 [36] 和标签平滑 [71]。作者所有的模型都是用AdamW [53; 41] 进行训练的。学习率缩放规则是 。在本文中，作者将批量大小设置为4096，所以学习率是 。作者的 MambaOut 模型是用 PyTorch [56] 和 timm [81] 库实现的，并在 TPU v3 上进行训练。更多的训练超参数在附录的表5中展示。

结果. 作者在表1中展示了 MambaOut 模型、视觉 Mamba 模型以及各种基于卷积和注意力的其他模型在 ImageNet [19; 66] 上的性能。值得注意的是，作者的 MambaOut 模型（不包含SSM）在所有模型大小上都一致地超过了包含SSM的视觉 Mamba 模型 [104; 50; 37; 58; 88]。例如，MambaOut-Small 模型达到了84.1%的top-1准确率，比 LocalVMamba-S [37] 高出0.4%，同时只需要79%的MACs。这些结果强烈支持作者的 _假设1_，即认为在 ImageNet 上进行图像分类时引入SSM是不必要的，符合奥卡姆剃刀原则。

此外，与最先进的基于卷积和注意力的模型相比，视觉 Mamba 模型目前表现出显著的性能差距。例如，采用简单可分离卷积 [67] 和标准注意力机制 [76] 的传统标记混合器的 CAFormer-M36 [92]，其准确率超过了所有相似大小的视觉 Mamba 模型超过1%。如果未来的研究旨在挑战作者的 _假设1_，将有必要### 对 COCO 上的目标检测和实例分割

设置. COCO 2017 [48] 作为广为人知的目标检测和实例分割基准。在作者的实验中，MambaOut 被用作 Mask R-CNN [31] 的 Backbone 网络，使用在 ImageNet 上预训练的权重进行初始化。作者遵循标准的1 训练计划，共12个周期。训练图像被调整大小，使得较短的一边为800像素，较长的一边不超过1333像素。使用 AdamW 优化器 [53, 41]，学习率为0.0001，总批量大小为16。作者的实现利用了 PyTorch [56] 和 mmdetection [8] 库。作者使用 FP16 精度来节省训练成本。实验在4个 NVIDIA 4090 GPU 上进行。

结果。尽管MambaOut在COCO[48]的目标检测和实例分割任务上可以超越一些视觉Mamba模型[58, 88]，但它仍然落后于最先进的视觉Mamba模型，如VMamba[50]和LocalVMamba[50]。例如，以MambaOut-Tiny作为Mask R-CNN的主干网络，其性能比VMamba-T[50]低1.4 和1.1 。这种性能差异强调了在长序列视觉任务中整合Mamba的好处，强化了作者的假设2。然而，在与其他最先进的卷积-注意力-混合模型，如TransNeXt[69]相比时，视觉Mamba仍然表现出显著的性能差距。视觉Mamba需要通过在视觉检测任务中超越其他最先进的模型来进一步验证其有效性。### ADE20K上的语义分割

设置。广泛用于语义分割任务的ADE20K[103]基准测试涵盖了150个语义类别。它包括训练集中的20,000张图片和验证集中的2,000张图片。在实验中，作者将Mamba作为UperNet[85]的主干网络，并从ImageNet预训练权重进行初始化。训练使用AdamW优化器[41, 53]，学习率为0.0001，批量大小为16，进行160,000次迭代。作者的实现使用了PyTorch[56]和mmsegmentation[14]库。实验在四块NVIDIA 4090 GPU上进行，使用FP16精度以提高训练速度。

结果。 在ADE20K上进行语义分割的性能趋势与在COCO上进行目标检测的趋势相似。MambaOut可以超越一些视觉Mamba模型，但无法匹敌最先进的Mamba模型的结果。例如，LocalVMamba-T[37]在单尺度(SS)和多尺度(MS)评估中均比MambaOut-Tiny高出0.5 mIoU，这进一步从实证上证实了作者的假设2。此外，视觉Mamba模型在与集成卷积和注意力机制的更先进混合模型，如SG-Former[65]和TransNeXt[69]进行比较时，继续表现出明显的性能不足。视觉Mamba需要进一步展示其在长序列建模方面的优势，通过在视觉分割任务中提供更强的性能。结论。

在本文中，作者从概念上讨论了Mamba机制，并得出结论它非常适合具有长序列和自回归特性的任务。作者根据这些标准分析了常见的视觉任务，并认为在ImageNet图像分类中引入Mamba是不必要的，因为它不符合这些特性。然而，对于符合长序列特性的视觉检测和分割任务，Mamba的潜力值得进一步探索。为了实证支持作者的观点，作者开发了MambaOut模型，这些模型采用了没有核心标记混合器SSM的Mamba块。MambaOut在ImageNet上超越了所有视觉Mamba模型，然而与最先进的视觉Mamba模型相比，它表现出明显的性能差距，从而验证了作者的论断。由于计算资源的限制，本文仅验证了视觉任务中的Mamba概念。将来，作者可能会进一步探索Mamba和RNN概念，以及将RNN和Transformers整合到大型语言模型（LLMs）和大型多模态模型（LMMs）中。

参考

[1].MambaOut: Do We Really Need Mamba for Vision?.