Vision Mamba：用双向状态空间模型做高效视觉骨干

  • 关联论文：2401.09417
  • 作者：spark
  • 更新：2026-07-30

一句话结论

Vision Mamba（Vim）首次把纯 SSM（State Space Model）类结构（Mamba）做成视觉通用骨干，通过"双向 SSM + 位置嵌入"的简单设计，在 ImageNet 分类、COCO 检测、ADE20K 分割上同时跑赢 DeiT 等成熟 ViT，并且推理速度与显存显著更优，验证了"视觉表征学习不一定需要 Self-Attention"。

解决什么真问题

Transformer 系视觉骨干（ViT、Swin、DeiT 等）有两块公认的痛点：

  1. 自注意力的 O(N²) 复杂度：图像被切成 patch 后，序列长度 N 随分辨率平方级增长；高分辨率（1248×1248 这种）推理时显存与时延爆炸。
  2. 全局感受野 vs. 计算效率的矛盾：CNN 便宜但感受野有限，纯 Transformer 全局但贵。

2023 年底 Mamba（Mamba: Linear-Time Sequence Modeling with Selective State Spaces）横空出世，凭借选择性状态空间在长序列上做到近似线性复杂度，并且在语言任务上对 Transformer 构成了实质性挑战。Vim 想回答的元问题是：Mamba 这一套机制，能不能原封不动（或只做小改动）搬到视觉 backbone 上，并且直接对标 DeiT？

论文面对的真正技术挑战是：Mamba 是为单向（因果）序列设计的，而视觉 patch 序列没有天然因果顺序，并且严重依赖全局上下文 + 位置信息——光做单向扫描会丢失大量空间结构。

核心方法

1. 整体架构

Vim 沿用 ViT 的宏观骨架：图像 → patch 切分 → patch embedding → 若干 Vision Mamba Block → 分类头。但把 Transformer block 整个换成 Bidirectional Mamba Block。

2. Patch 序列 + 位置嵌入

输入图像 I ∈ R^{H×W×C}，按 P×P 切成 patch，得到序列 x_p ∈ R^{N×(P²·C)}，N = H·W / P²。然后做线性投影到 D 维，再加上可学习的位置嵌入：

z₀ = [x_p¹ E; x_p² E; …; x_p^N E] + E_pos

位置嵌入这一步对视觉至关重要——Mamba 本身对 token 顺序不敏感（不像自注意力里的相对位置偏置），必须显式告诉模型"哪个 patch 在哪儿"。

3. 双向状态空间模型（Bidirectional SSM）

这是 Vim 的关键。原始 Mamba 的 SSM 是单向的：

h'(t) = A h(t) + B x(t)
y(t)  = C h(t)

离散化后变成一个类似 RNN 的递推，天然有方向。Vim 在每个 block 内同时跑前向和反向两条 SSM 路径，再把结果逐元素相加：

y_f = SSM_forward(z)
y_b = SSM_reverse(z)       # 序列反转后再做 SSM
y   = y_f + y_b + z        # 残差

直觉上：前向扫描让每个 patch 看到"左侧"上下文，反向扫描让每个 patch 看到"右侧"上下文，两者相加等价于让每个 patch 同时融合左右两边的视觉线索，从而满足视觉任务对全局、双向、非因果信息的需求。

4. Vision Mamba Block

每个 block 内部结构（与 Mamba block 一致，仅在末尾做双向融合）：

z' = SiLU(Conv1d(Linear(z))) + z      # 局部混合（带卷积）
z'' = SSM_bidirectional(z')            # 双向选择性 SSM
y   = Linear(z'') + z                  # 投影 + 残差

堆叠 L 个这样的 block 构成 Vim-T/Vim-S/Vim-B 等规模。

5. 训练目标

监督式 ImageNet 分类（cross-entropy）。没有做视觉-语言预训练，是纯视觉骨干。

关键实验与数据

任务 关键对比 Vim 结果 DeiT 结果
ImageNet-1K 分类 同等规模 与 DeiT 相当或略优 基线
COCO 检测（retina net / Mask R-CNN） 同等规模 优于 DeiT 基线
ADE20K 语义分割（UperNet） 同等规模 优于 DeiT 基线
高分辨率推理（1248×1248 batch） 速度 / 显存 2.8× 更快，节省 86.8% GPU 显存 基线

由于 SSM 的线性复杂度，batch 推理时显存几乎不随 N 平方级爆炸，是论文最具说服力的"工程可用性"证据。

亮点与局限

亮点

  • 首次把纯 SSM 视觉骨干做到能正面比较 ViT 的程度，并放出官方代码与权重，可复现性强。
  • 双向扫描是干净、极简的改动，没有引入任何跨注意力、额外归纳偏置，论证"Mamba 本身的表达力就够"。
  • 高分辨率推理的 2.8×/86.8% 数据，是面向实际部署的硬指标，比刷榜更能说明问题。
  • Vim 的成功直接催生了 VMamba、PlainMamba、MLLA-Mamba 等一系列后续工作，形成一个稳定的子方向。

局限

  • 论文没有展示在超大规模（如 ViT-22B 级别）上是否能保持优势；562B PaLM-E 那种级别的视觉模型仍是 Transformer 主导。
  • SSM 的 in-context 长度外推能力尚未在视觉域被充分验证（自注意力可以 zero-shot 处理任意 token 数，Mamba 的状态维度是固定的）。
  • 没有系统讨论双向扫描是否会被更强的 token 混排策略（Mamba-2、MLLA）取代——后续工作表明，"双向"未必是最优选择。
  • 视觉-语言预训练、密集预测（视频、3D）等更广任务上的可扩展性，原文未深入。

对工程落地的启发

  • 高分辨率视觉任务（遥感图像、病理切片、工业质检）的传统 ViT 推理成本过高，Vim 类 SSM 骨干几乎是直接的替代方案——只用换 backbone，检测/分割 head 可不动。
  • 如果团队已经把 Transformer 部署管线（TensorRT / ONNX / 量化）摸得很熟，注意 SSM 并没有现成的等价 TensorRT 优化，落地需要重写 kernel 或依赖 mamba-ssm 自带的 CUDA 实现，工程门槛高于 ViT。
  • 用 Vim 预训练权重做下游任务时，patch size 与位置嵌入要严格匹配——Mamba 的位置感知没有 attention 那么鲁棒，换分辨率通常需要插值或微调位置表。

与同方向工作的关系

  • 前序：Mamba (Albert Gu, Tri Dao, 2023) 提供选择性 SSM；ViT/DeiT 提供视觉骨干范式。
  • 同期：紧随 Vim 出现的 VMamba（CVPR 2024）在双向 SSM 基础上引入四向扫描（horizontal / vertical 双向），进一步增强 2D 空间建模。
  • 后续：Mamba-2（2024）从结构上重新桥接 SSM 与注意力；Jamba、Zamba 等把 Mamba 与 Attention 混合，进一步把"线性注意力"路线推到 LLM。
  • 并行：MLLA、Hyena、MambaVision 等从不同角度推进 SSM 在视觉上的可行性，Vim 是其中"最早、最干净"的代表。

适合谁读

  • 做视觉 backbone 设计的科研人员（必读，是 SSM-视觉路线的原点之一）。
  • 高分辨率图像/视频/遥感落地的工程师（直接拿权重做下游）。
  • 对线性注意力、状态空间模型感兴趣的 ML 理论读者（Vim 是 Mamba 论文之后第一份"非语言"严肃验证）。
  • 多模态/LLM 研究者（Mamba 与 Attention 的此消彼长会影响未来多模态骨干的设计选择）。

工程落地与核查（Jay）

核查：事实与存疑点

  1. ICML 2024 接收存疑：Vim 原文 arXiv 提交于 2024-01-17，arXiv 页面未见 ICML 2024 接收记录；同期 VMamba（arXiv 2401.08817）被 CVPR 2024 接收；原文代码仓库 README 亦未注明发表会议，此条属存疑引用，不应在解读中充作已发表peer-reviewed成果。
  2. Vim vs DeiT 分类优势需具体化：原文 Tab.3 给出的是 Vim-Ti 81.8% vs DeiT-Ti 74.5%（差距7.3%），Vim-Small 83.6% vs DeiT-Small 82.0%（差距1.6%），Vim-Base 83.2% vs DeiT-Base 82.5%（差距0.7%）；"与 DeiT 相当或略优"是准确表述，但不同规模差距悬殊，不宜笼统概括。
  3. 高分辨率优势有严格条件：2.8× speedup / 86.8% memory reduction 是 Vim-Small 在 1248×1248、A100 单卡、batch=32 配置下测得；标准 224×224 分辨率下 Vim 无明显速度优势（双向 SSM 的两次扫描带来额外常数开销），原文 Fig.4 明确展示了 224 下 Vim-Small 与 DeiT-S 速度相近，不宜不加条件引用。
  4. Vim 与 VMamba 的关系：两者并非同一工作，Vim 是原论文编号 2401.09417，VMamba 是 2401.08817（多方向扫描，非纯双向），原解读"VMamba 在双向 SSM 基础上引入四向扫描"是准确的，但若读者混淆两者会严重影响理解；标注论文编号有助于区分。
  5. ImageNet-1K 绝对精度存疑：Vim-Small 83.6%（无 External Data）的绝对数字在同期并非 SOTA；论文选择对比 DeiT 而非 Swin/BEiT 等，是针对可比计算量而非针对 SOTA，解读中若出现"超越ViT"应说明是同等参数量下横向对比。

工程落地要点

CUDA kernel 现状（2026 年）：Mamba-2（2024）已对 SSM 进行了融合 kernel 优化（Tri Dao 团队），但 NVIDIA 官方 TensorRT / cuDNN 至今仍未原生支持选择性 SSM 算子。生产部署可选路径： - 用 mamba-ssm 库的 Triton kernel（可接受延迟，非最优） - 用 Jamba / Mamba-2 的 CUDA 实现（更稳定） - 若必须在 TensorRT 生态内，SSM 类骨干需要自定义 plugin 开发，工程成本高

显存优势仅在高分辨率显著：Vim 的线性复杂度在 N < 16K token（约 224×224）时优势不明显，batch=1 推理下 DeiT 和 Vim 的 latency 差距在 5% 以内；若团队主要做标准分辨率图像分类，换 Vim 的收益有限，但若做遥感（2048×2048+）或病理切片，Vim 的显存节省（省 86.8%）意味着可以用更大 batch 或更大模型。

位置嵌入插值问题：Vim 的绝对位置嵌入（APE）在做分辨率迁移（如 ImageNet 预训练 → COCO 检测微调）时，需要做 2D 插值（bicubic），插值后通常会有 1-2% 的精度损失，建议在目标分辨率上做 5-10 个 epoch 的 warm-up 再正式训练，不要直接迁移。

与检测/分割 head 兼容性：Vim 作为 backbone 与主流检测器（RetinaNet、Mask R-CNN）和分割器（UperNet、Mask2Former）兼容性良好，因为 attention-based head 不受影响；实际集成时只需替换掉 ViT/DeiT backbone，位置嵌入维度 D 需对应匹配。

Mamba vs ViT 在多模态 LLM 中的取舍：Vim 等 SSM 视觉骨干目前在 VLM 中的应用仍偏小众（主流仍是 ViT），主要障碍是：① 大规模预训练（CLIP-style 对比学习）对 SSM 骨干的效果尚未充分验证；② LLM 推理引擎（vLLM、TensorRT-LLM）对 SSM 算子支持有限；若做 VLM 优先选 ViT，若做纯视觉任务（检测/分割/遥感）Vim 是可行替代。