论文标题：MambaOut: Do We Really Need Mamba for Vision?
论文地址：https://arxiv.org/pdf/2405.07992
摘要
Mamba 这种基于状态空间模型（SSM）的架构，虽然可以解决自注意力机制的平方复杂度问题，但在视觉任务中的表现并不如顶尖的卷积或注意力机制的模型。
Mamba 更适合具有长序列和自回归特性的任务，但图像分类不符合这两点，所以 Mamba 对于图像分类并不必要。
然而，检测和分割任务符合长序列特性，虽然不是自回归的，因此仍然有探索 Mamba 在这些任务中潜力的价值。
MambaOut 系列模型通过去掉核心的 SSM 模块，验证了 Mamba 在图像分类任务中的不必要性，但对于检测和分割任务，Mamba 仍然表现出较强的潜力
提出假设
Mamba 是否适合视觉任务？文章分析了 Mamba 的机制，指出其适合具有长序列和自回归特性的任务。图像分类任务既不具有长序列特性，也不自回归；检测和分割任务虽然符合长序列特性，但不是自回归的。
本文提出两种假设：
- 假设1：SSM 对于图像分类并不必要，因为它不符合长序列和自回归特性。
- 假设2：SSM 可能对目标检测和分割任务有帮助，因为它们符合长序列特性，尽管不是自回归的。
通过实验表明假设都是正确的。
mamba 适合什么任务？
长序列且自回归的任务。
类似于RNN的SSM的本质是因果注意力模型，隐藏状态h可以被视为一个固定大小的记忆卡，存储历史信息，通过更新记忆来更新记忆卡，卡的大小是恒定的。所以可以说SSM本质上记忆是有损的，相比于Transformer，短序列上表现的并不好，但是在长序列上的表现要优于Transformer。
上图(a)展示的是因果注意力模型，每一步都会保存历史的k和v，随着序列变成，计算复杂度也会增加。
(b)展示的是类RNN的模型（SSM），将过去状态压缩到固定大小的hidden state中，记忆大小固定，所以记忆会丢失，计算复杂度不变，更适合长序列。
觉图片识别任务上，并不需要非常长的序列。
在这个分析中，输入的大小为 X ∈ R L × D X \in \mathbb{R}^{L \times D} X∈RL×D，其中 L L L 是 token 长度（即序列长度）， D D D 是通道或嵌入的维度。根据这个设定，FLOPs 可以计算为：
F L O P s = 24 D 2 L + 4 D L 2 \mathrm{FLOPs} = 24D^2L + 4DL^2 FLOPs=24D2L+4DL2这里，第一项表示线性复杂度，第二项是与 L L L 的二次复杂度。
为了比较二次项和线性项的计算负担，引入了一个比率 r L r_L rL： r L = 4 D L 2 24 D 2 L = L 6 D r_L = \frac{4DL^2}{24D^2L} = \frac{L}{6D} rL=24D2L4DL2=6DL
当 L > 6 D L > 6D L>6D 时，二次复杂度的计算负担将超过线性复杂度，这意味着该任务具有长序列特征。
举例来说：
对于 ViT-Small（ViT-S），通道数 D = 384 D = 384 D=384，因此阈值为：
τ small = 6 × 384 = 2304 \tau_{\text{small}} = 6 \times 384 = 2304 τsmall=6×384=2304对于 ViT-Base（ViT-B），通道数 D = 768 D = 768 D=768，因此阈值为：
τ base = 6 × 768 = 4608 \tau_{\text{base}} = 6 \times 768 = 4608 τbase=6×768=4608对于 ImageNet 上的图像分类任务，典型的输入图像大小为 22 4 2 224^2 2242，这意味着每个图像分割为 1 6 2 16^2 162 的 patch 后生成 1 4 2 = 196 14^2 = 196 142=196 个 token。显然，196 远小于 τ small \tau_{\text{small}} τsmall 和 τ base \tau_{\text{base}} τbase，这表明 ImageNet 上的图像分类任务并不符合长序列任务的特征。
对于 COCO 数据集上的目标检测和实例分割任务，推理图像的大小为 800 × 1280 800 \times 1280 800×1280，而 ADE20K 数据集上的语义分割推理图像的大小为 512 × 2048 512 \times 2048 512×2048。这些情况下，token 数量大约为 4K。由于 4 K > τ small 4K > \tau_{\text{small}} 4K>τsmall 且 4 K ≈ τ base 4K \approx \tau_{\text{base}} 4K≈τbase，因此这两个任务可以被视为长序列任务。
总的来说，这一分析表明，图像分类任务通常不是长序列任务，而目标检测和语义分割等任务由于图像的高分辨率，具有长序列的特征。
全可见与因果模式的适用性： 在视觉识别任务中，全可见模式允许模型一次性访问所有输入信息，因此对理解任务（如图像分类、物体检测和语义分割）更为合适。而因果模式限制了每个 token 只能访问之前的 token，因此更适合自回归任务（如语言生成）。
因果限制对视觉任务的影响： 对视觉任务施加因果限制会导致性能下降，尤其是对于像 Vision Transformers 这样依赖全局信息的模型。视觉任务不需要因果 token 混合，因为它们并不具有自回归特性。
实验证明
本文设计了实验来佐证假设：
根据图 1(a) 所示，Mamba 模块是基于 Gated CNN 模块构建的。Gated CNN 和 Mamba 的元架构可以看作是 MetaFormer 的 token mixer 和 MLP 的简化集成，类似于 MetaNeXt 。给定输入 X ∈ R N × D X \in \mathbb{R}^{N \times D} X∈RN×D，其元架构公式如下：
X ′ = N o r m ( X ) , Y = ( T o k e n M i x e r ( X ′ W 1 ) ⊙ σ ( X ′ W 2 ) ) W 3 + X , X' = \mathrm{Norm}(X), \\ Y = (\mathrm{TokenMixer}(X'W_1) \odot \sigma(X'W_2))W_3 + X, X′=Norm(X),Y=(TokenMixer(X′W1)⊙σ(X′W2))W3+X,
其中， N o r m ( ⋅ ) \mathrm{Norm}(\cdot) Norm(⋅) 表示归一化操作， T o k e n M i x e r ( ⋅ ) \mathrm{TokenMixer}(\cdot) TokenMixer(⋅) 是用于进行 token 混合的模块， W 1 ∈ R D × r D W_1 \in \mathbb{R}^{D \times rD} W1∈RD×rD， W 2 ∈ R D × r D W_2 \in \mathbb{R}^{D \times rD} W2∈RD×rD 和 W 3 ∈ R r D × D W_3 \in \mathbb{R}^{rD \times D} W3∈RrD×D 是带有 MLP 扩展比 r r r 的可学习参数， σ \sigma σ 是激活函数。Gated CNN 和 Mamba 的 token mixer 分别为：
T o k e n M i x e r G a t e d C N N ( Z ) = C o n v ( Z ) , T o k e n M i x e r M a m b a ( Z ) = S S M ( σ ( C o n v ( Z ) ) ) . \mathrm{TokenMixer}{\mathrm{Gated\,CNN}}(Z) = \mathrm{Conv}(Z), \\ \mathrm{TokenMixer}{\mathrm{Mamba}}(Z) = \mathrm{SSM}(\sigma(\mathrm{Conv}(Z))). TokenMixerGatedCNN(Z)=Conv(Z),TokenMixerMamba(Z)=SSM(σ(Conv(Z))).
通过比较公式，以及参考图 1(a)，Gated CNN 和 Mamba 模块 之间的主要区别在于是否包含 SSM。因此，本文开发了一系列名为 MambaOut 的模型，它基于 Gated CNN 模块而没有使用 SSM。MambaOut 评估 Mamba 在视觉识别任务中的必要性。
具体而言，我们将 Gated CNN 的 token mixer 指定为 7 × 7 核大小的 depthwise 卷积 ，参考 ConvNeXt。此外，为了提高实际速度，我们仅在部分通道上执行 depthwise 卷积 ，参考 InceptionNeXt。如算法 1 所示
Gated CNN 模块的实现既简单又优雅。类似于 ResNet，我们通过在每个阶段堆叠 Gated CNN 模块构建了 MambaOut，如图 4 所示。
每个模型尺寸的配置细节显示在附录中的表 4。
实验结果
1. 图像分类 (ImageNet)
MambaOut在ImageNet图像分类任务中表现优异，超过了包含SSM的Mamba模型。具体结果如下：
MambaOut-Small 模型的Top-1准确率为 84.1%，比 LocalVMamba-S 提高了 0.4%，同时只需要 79% 的MACs计算量。
实验结果表明，MambaOut模型在没有SSM的情况下，仍然能够在ImageNet上取得优异的分类性能，验证了假设1：在ImageNet图像分类任务中，引入SSM是不必要的。
2. 目标检测和实例分割 (COCO)
在COCO数据集上，MambaOut作为Mask R-CNN的骨干网络进行目标检测和实例分割任务。结果表明：
MambaOut-Tiny 在目标检测任务中比 VMamba-T 落后了 1.4 APb 和 1.1 APm，表明在长序列视觉任务中，Mamba的优势仍然存在。
虽然MambaOut在某些情况下超过了部分Mamba模型，但在与最先进的视觉Mamba模型（如 VMamba 和 LocalVMamba）的对比中，仍有性能差距。这验证了假设2：尽管Mamba在图像分类任务中表现不佳，但在长序列视觉任务（如目标检测和分割）中仍然值得进一步探索。
3. 语义分割 (ADE20K)
在ADE20K语义分割任务中，MambaOut与现有的Mamba模型进行对比：
MambaOut-Tiny 模型在单尺度 (SS) 和多尺度 (MS) 评估中的mIoU值比 LocalVMamba-T 分别低了 0.5。
结果表明，MambaOut虽然在某些情况下超越了一些Mamba模型，但在最先进的Mamba模型上，MambaOut依然表现稍逊，验证了Mamba在长序列建模方面的潜力。
总结：
MambaOut在ImageNet分类任务中，成功验证了SSM并不必要，且能有效超越现有的视觉Mamba模型。
在目标检测和语义分割等长序列任务中，Mamba仍然具有优势，MambaOut无法完全达到最先进的Mamba模型的性能。