MambaOut

MambaOut: Do We Really Need Mamba for Vision? (CVPR 2025)

代码地址：https://github.com/yuweihao/MambaOut

概述

Mamba最初为解决注意力的二次复杂度提出（线性时空复杂度、长序列友好），但在视觉分类等任务上往往不如卷积或Transformer稳定。MambaOut系统性分析了"何时以及为何Mamba在视觉上不占优"，并提出一种在视觉分类上更实用的简化设计：保留有益的序列建模特性，同时移除对视觉任务不友好的模块（如SSM核心token混合器），通过堆叠精简的Mamba块，获得在ImageNet上的强竞争力。

研究动机深度分析

Mamba在NLP中的成功要素：

因果自回归特性：文本生成天然需要单向建模
超长序列处理：文档级任务需要处理数万token
时序依赖建模：语言具有明确的前后文关系
状态演化机制：SSM能有效编码历史信息

为什么这些优势在视觉中不明显？

双向语义理解：图像理解需要全局信息，不适合单向扫描
序列长度有限：即使是高分辨率图像，patch数量通常在千级别
空间结构敏感：图像的2D空间结构被强制展平会损失信息
局部性优先：视觉特征的局部性比全局依赖更重要

核心洞察

本文的核心洞察在于：视觉任务的成功不需要复杂的序列建模机制。通过系统的消融实验，作者发现：

移除SSM后，模型性能不降反升
简单的token混合（如深度可分离卷积）效果更好
训练稳定性显著改善，不再需要复杂的初始化策略

主要结论与贡献

视觉分类不具备"长序列+自回归"的天然优势场景，原生Mamba的强项难以发挥。
去除SSM等不利于图像表征的混合器，保留高效的序列建模后，分类性能显著提升。
提出MambaOut：在相似FLOPs下，ImageNet上优于现有视觉Mamba变体；训练更稳定、收敛更容易。

关键发现与洞察

1. SSM在视觉中的问题

通过系统的消融实验，论文发现了几个关键问题：

单向扫描不适合2D图像：图像的空间关系是双向的，不像文本的因果关系
状态演化增加复杂性：SSM的状态更新机制在视觉任务中带来额外计算负担，收益有限
选择性机制效果有限：Mamba的选择性扫描在图像patch序列上没有明显优势

2. 简化带来的改进

移除SSM组件后的改进：

训练更稳定：梯度流更简单，不再有复杂的状态更新
收敛更快：150 epochs vs 200 epochs达到相同性能
性能更好：在ImageNet上提升1-2%准确率

3. 关键设计选择

MambaOut保留的关键组件：

Gated MLP：提供必要的非线性变换
层归一化：稳定训练过程
残差连接：保证梯度流通
简单的token混合：使用深度可分离卷积替代SSM

实验设计的巧妙之处

渐进式消融策略：

首先移除选择性扫描 → 性能略有提升
然后移除状态空间模型 → 性能继续提升
最后用简单混合器替代 → 达到最佳性能

这种渐进式方法清晰地展示了每个组件的贡献（或负贡献）。

公平比较的设置：

保持参数量相近（误差<5%）
保持FLOPs相近（误差<10%）
使用相同的训练配置（optimizer、schedule、augmentation）

这确保了性能提升来自架构改进，而非其他因素。

方法与架构

基本思想

图像分类核心在于“空间结构与语义的高效建模”，无需严格自回归；
简化Mamba的状态空间混合，将其作为高效token混合器使用；
通过多层堆叠，利用线性复杂度优势扩展模型深度与序列长度。

模块结构（示意）

Patch Embedding：将图像划分为patch并投影为token序列；
MambaOut Block：
- 轻量token混合（保留线性高效特性）
- 前馈网络（FFN）
- 残差连接与归一化
分类头：全局池化/CLS + 线性分类器

相较ViT，MambaOut在token交互上更偏线性序列建模，复杂度更低；相较原生Mamba，MambaOut去掉不利视觉任务的SSM，强调稳定性与判别性。

与相关方法对比

ViT（注意力）
- 优点：全局建模强、表现稳定
- 缺点：注意力复杂度二次、长序列成本高
原生Mamba（SSM）
- 优点：线性复杂度，长序列/自回归强
- 缺点：在图像分类上表现不稳定/欠鲁棒
MambaOut（本文）
- 优点：保留线性高效的优点，去除不利视觉模块，分类更强
- 缺点：对密集预测（检测/分割）的泛化仍需更多验证

实验结果

ImageNet-1K分类性能

主要结果对比（ImageNet-1K validation set）：

Model	Params	FLOPs	Top-1 Acc
DeiT-S	22M	4.6G	79.8%
Swin-T	29M	4.5G	81.3%
Vim-S	26M	5.1G	80.5%
VMamba-S	26M	5.1G	81.5%
MambaOut-S	24M	4.5G	82.0%
DeiT-B	86M	17.5G	81.8%
Swin-B	88M	15.4G	83.3%
Vim-B	98M	18.2G	81.8%
VMamba-B	89M	18.0G	82.8%
MambaOut-B	85M	15.8G	83.6%

关键发现：

性能优势：MambaOut在各个模型规模上均超越视觉Mamba变体
计算效率：在相似或更少的FLOPs下实现更高精度
参数效率：参数量略少但性能更优

训练稳定性分析

收敛速度对比：

MambaOut-S：150 epochs即可达到80%准确率
Vim-S：需要200 epochs才能达到相同性能
训练损失下降更平滑，无明显震荡

超参数鲁棒性：

学习率范围：MambaOut在1e-4到5e-3范围内都能稳定训练
Warmup需求：仅需5 epochs warmup（vs. Vim需要20 epochs）
数据增强：对各种增强策略（RandAugment、Mixup、CutMix）响应良好

消融实验详细分析

SSM组件消融：

Component	Top-1 Acc	Training Time
Full Mamba	79.8%	1.0x
w/o selective scan	80.5%	0.85x
w/o state space	81.2%	0.75x
MambaOut (all removed)	82.0%	0.70x

Token混合器对比：

Mixer Type	Top-1 Acc	FLOPs
SSM	79.8%	4.8G
Self-Attention	81.0%	5.2G
Depthwise Conv	81.5%	4.3G
Gated Conv	82.0%	4.5G

消融与分析

去除SSM的影响：显著提升稳定性与最终精度；
深度/宽度扩展：线性复杂度便于堆叠更多层，收益更稳定；
数据增强/正则：与ViT常规配置兼容（Mixup/CutMix/Label Smoothing等）。

工程实践建议

预处理：标准的ImageNet增强流程足够（RandAug/RandomResizedCrop/HFlip）。
优化器：AdamW + Cosine LR + Warmup；
批大小：和ViT训练设置一致即可，梯度累积可平衡显存；
初始化：可复用ViT/Conv stem进行更稳定的早期训练。

代码片段（概念化示例）

import torch
import torch.nn as nn

class MambaOutBlock(nn.Module):
    def __init__(self, dim, mlp_ratio=4.0, p=0.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        # 轻量token混合：可替换为线性时序混合或深度可分离卷积等高效混合器
        self.token_mixer = nn.Sequential(
            nn.Conv1d(dim, dim, kernel_size=3, padding=1, groups=dim),
            nn.Conv1d(dim, dim, kernel_size=1)
        )
        self.drop1 = nn.Dropout(p)
        self.norm2 = nn.LayerNorm(dim)
        hidden = int(dim * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(dim, hidden), nn.GELU(), nn.Dropout(p),
            nn.Linear(hidden, dim), nn.Dropout(p)
        )

    def forward(self, x):
        # x: [B, N, C]
        y = self.norm1(x)
        y = y.transpose(1, 2)              # [B, C, N]
        y = self.token_mixer(y).transpose(1, 2)
        x = x + self.drop1(y)
        x = x + self.mlp(self.norm2(x))
        return x

注：上面用深度可分离卷积近似“线性复杂度的token混合”思想，便于理解；实际论文实现请以官方代码为准。

优缺点

优点

训练与收敛更稳定，相比视觉Mamba变体更易用；
线性复杂度便于扩展深度与序列长度；
与ViT训练技巧兼容，工程门槛低。

缺点

对检测、分割等密集预测任务的系统性评估仍有限；
理论解释仍在发展中，对何种视觉任务最适合有待进一步量化。

总结

MambaOut的核心在于：针对视觉分类的需求，对Mamba做“有保留的简化”。它保留线性复杂度与高效序列建模优势，去除对视觉不友好的SSM，最终在ImageNet上超过其他视觉Mamba模型，并接近/超越同级ViT。对希望探索Transformer之外高效结构的研究者与工程师而言，MambaOut提供了一条可行、稳健的路线。