MambaOut: 在视觉中我们真的需要Mamba吗？

0. TL;DR

Mamba 是一种具有状态空间模型（SSM）的 RNN 风格的 token mixer，旨在解决注意力机制的二次复杂度问题。然而，在视觉任务中，Mamba 的表现往往不如卷积和注意力模型。本文通过对 Mamba 的分析，提出了两个假设：

  1. 在图像分类任务中，由于不满足长序列和自回归特性，Mamba 并非必要；
  2. 而在目标检测和语义分割任务中，由于符合长序列特性，Mamba 仍具潜力。

通过构建 MambaOut 模型（移除了 SSM 的 Mamba 块），实验证明在 ImageNet 图像分类任务中，MambaOut 超越了视觉 Mamba 模型，验证了假设 1；在检测和分割任务中，MambaOut 未能匹配最先进的视觉 Mamba 模型，进一步支持了假设 2。

1. 背景介绍

在计算机视觉领域，Transformer 架构因其在长序列任务中的表现而备受关注。然而，传统的 Transformer 架构在视觉任务中面临计算复杂度的挑战，尤其是其注意力机制对序列长度存在二次复杂度的依赖。随着视觉任务对模型处理长序列的能力要求越来越高，研究人员开始探索替代方案。

Mamba 架构以其 RNN 风格的 token mixer 和状态空间模型（SSM）为核心，被提出用于解决注意力机制的复杂度问题。它旨在以线性复杂度处理长序列任务，尽管如此，在视觉任务中的表现却不如人意。

2. MambaOut 模型

Mamba 块基于 Gated CNN 块，两者在架构上类似，但关键的区别在于 Mamba 块引入了 SSM。Gated CNN 块主要由深度可分离卷积和门控机制组成，而 Mamba 块在此基础上添加了 SSM，以增强其在长序列任务中的表现。

SSM 是 Mamba 的核心组成部分，通过维护一个固定大小的隐藏状态来存储历史信息。其序列到序列的变换可表示为：

\[\begin{aligned} \mathbf{x}^\prime(t) &= A \mathbf{x}(t) + B \mathbf{u}(t) \\ \mathbf{y}(t) &= C \mathbf{x}(t) \\ \end{aligned}\]

为了验证 Mamba 在视觉任务中的必要性，本文提出了一种基于 Gated CNN 块构建的 MambaOut 模型。该模型移除了 Mamba 块中的 SSM，以评估其在图像分类任务中的性能。

MambaOut 模型采用分层架构，包含四个阶段的 Gated CNN 块。每个块由深度可分离卷积和门控机制组成，通道维度和块的数量根据模型大小进行调整。

Gated CNN 块的前向传播公式为：

\[X^\prime = Norm(X) \\ Y = (TokenMixer(X^\prime W_1) \odot \sigma(X^\prime W_2)) W_3 + X\]

3. 实验分析

在 ImageNet 图像分类数据集上，使用标准的训练方案，包括随机裁剪、水平翻转、RandAugment 等数据增强技术。模型使用 AdamW 优化器，学习率设置为 0.004，训练 300 个 epoch。

MambaOut 模型在 ImageNet 图像分类任务中表现出色，超越了所有视觉 Mamba 模型。例如，MambaOut-Small 模型的 Top-1 准确率为 84.1%，而 LocalVMamba-S 模型的准确率为 83.7%。这些结果支持了假设 1，即在图像分类任务中，Mamba 并非必要。

在 COCO 目标检测和实例分割数据集上，使用 Mask R-CNN 作为评估框架，训练 12 个 epoch。模型在图像尺寸为 800×1280 的情况下进行评估。

尽管 MambaOut 模型在某些情况下超过了部分视觉 Mamba 模型，但仍未能达到最先进的视觉 Mamba 模型的性能。例如，MambaOut-Tiny 模型的 APb 为 45.1，而 VMamba-T 模型的 APb 为 46.5。这些结果验证了假设 2，即在长序列任务中，Mamba 仍具潜力。

在 ADE20K 语义分割数据集上，使用 UperNet 作为评估框架，训练 160,000 次迭代。模型在图像尺寸为 512×2048 的情况下进行评估。

MambaOut 模型在语义分割任务中的表现与目标检测任务类似，未能达到最先进的视觉 Mamba 模型的性能。例如，MambaOut-Tiny 模型的 mIoU 为 47.4，而 LocalVMamba-T 模型的 mIoU 为 47.9。这些结果进一步支持了假设 2。