在过去几年里，视觉 Transformer (ViT) 一直主导着计算机视觉领域。通过将图像视为图像块序列并应用自注意力机制，这些模型在图像分类、目标检测和语义分割方面树立了新的标杆。然而，这种强大的能力是以高昂的计算成本为代价的。

Transformer 核心的自注意力机制存在二次方复杂度的问题。简单来说，如果你将图像块的数量增加一倍 (例如，通过提高分辨率) ，计算量和内存需求不仅仅是翻倍，而是翻四倍。这使得高分辨率图像处理变得缓慢、消耗大量内存，并且在没有专用硬件或繁琐架构优化的情况下往往难以实现。

但如果我们能在保留 Transformer 的全局上下文感知能力的同时，又避免二次方复杂度的瓶颈，那会怎样呢？这个问题引导研究人员探索替代方案。一个很有前途的候选方案来自一个意想不到的领域:** 经典控制理论**。状态空间模型 (SSM) 就此登场，以及其最新的强大化身——Mamba。在自然语言处理领域，Mamba 凭借其随序列长度线性扩展的能力，已经达到了甚至在某些情况下超越了 Transformer 的水平。

一篇新论文《Vision Mamba: 基于双向状态空间模型的高效视觉表示学习》迈出了合乎逻辑的下一步: 将 Mamba 应用于计算机视觉。所提出的 Vision Mamba (Vim) 骨干网络使用 SSM 处理图像，达到了 Transformer 级别的性能——甚至更好——同时效率显著提高。

DeiT (视觉 Transformer) 与 Vim (Vision Mamba) 的性能比较摘要。Vim 在分类、分割和检测任务中准确率更高，同时速度显著更快，并且在高分辨率图像下使用的 GPU 显存大幅减少。

图 1: Vim 与 DeiT 视觉 Transformer 的正面对比。Vim 在各项任务的准确率上持续优于 DeiT，同时在速度和内存效率方面表现更佳，尤其是在高分辨率下。

在本文中，我们将深入解读 Vision Mamba 这篇论文，介绍状态空间模型的基础知识，了解作者如何将其应用于视觉领域，并审视这些可能预示着计算机视觉骨干网络设计转变的惊人成果。

从 Transformer 到状态空间模型

为了理解 Vim 的影响力，让我们从它要解决的问题开始。

视觉 Transformer (ViT) 的工作方式是将图像切成多个图像块 (例如 16×16 像素) ，将它们展平为向量，并将这些向量视为一个序列——类似于句子中的单词。自注意力机制会计算每对图像块之间的相互作用，使模型能够捕捉全局上下文。例如，猫耳朵的一个图像块可以“关注”到它尾巴的一个图像块，无论它们之间的空间距离有多远。缺点是: 计算长度为 N 的序列中每对元素的相互作用会导致 O(N²) 的复杂度。

状态空间模型 (SSM) 采用了不同的方法。受连续系统的启发，它们定义了一个潜在“状态” \( h(t) \)，其随时间演变，遵循以下公式:

\[ h'(t) = A h(t) + B x(t) \]

\[ y(t) = C h(t) \]

这里，\( A, B, C \) 是描述系统动态的参数。为了将 SSM 应用于深度学习，我们使用一个步长 \( \Delta \) 对这些方程进行离散化，得到离散参数 \( \overline{A} \) 和 \( \overline{B} \):

用于离散化连续状态空间参数 A 和 B 的方程。

离散化后，该模型可以用两种形式表示:

  • 循环形式: 逐步处理序列，在每个元素处更新隐藏状态。 离散状态空间模型的循环形式。

  • 卷积形式: 将循环过程表示为对整个序列的一次大型卷积，从而实现并行训练。 状态空间模型的卷积形式，其中输出 y 是输入 x 与一个结构化卷积核 K 之间的卷积结果。

Mamba 架构的改进在于使参数 \( B, C, \Delta \) 依赖数据——由当前输入序列决定——从而实现类似注意力的选择性适应，并保持线性复杂度。

为视觉领域改造 Mamba: Vim 架构

将为一维文本序列构建的 Mamba 转化为图像处理模型，需要巧妙的工程设计。

Vision Mamba (Vim) 架构概览。输入图像被分割成图像块，嵌入位置信息，然后由使用双向状态空间模型的 Vim 编码器块进行处理。

图 2: Vision Mamba 遵循类似 ViT 的“分块-嵌入”流程，但将自注意力替换为基于双向 SSM 的 Vim 块。

Vim 的处理流程:

  1. 分块: 提取不重叠的图像块。
  2. 嵌入: 将每个图像块展平并通过线性投影转化为一个令牌 (token)。位置嵌入提供空间上下文，也可以加入分类令牌。 该方程示例展示了图像块如何被投影并与位置嵌入结合，形成初始令牌序列。
  3. Vim 编码器: 令牌通过一个包含 \( L \) 个双向 Vim 块的堆栈。
  4. 预测: 将分类令牌的输出送入 MLP 头进行预测。

双向 Vim 块: 从所有方向捕捉上下文

标准的 Mamba 块只能单向处理序列，这适用于文本，但不适用于图像。像素需要来自各个相邻方向的上下文。Vim 通过双向处理解决了这一问题:

在每个 Vim 块中:

  1. 输入归一化。
  2. 计算两个投影:
    • \( x \): 主分支特征。
    • \( z \): 门控向量。
  3. 前向分支: 使用 Conv1d 和 SSM 从头到尾处理图像块。
    后向分支: 反向处理图像块。
  4. 每个分支通过各自的选择性 SSM 捕捉对应方向的上下文。
  5. 两个分支的输出与 \( z \) 进行门控并相加。
  6. 最后的线性投影和残差连接完成该块。

通过融合前向和后向的上下文，Vim 构建了类似自注意力的全局表示——但避免了二次方复杂度的开销。

Vim 为何高效

三大优化使 Vim 脱颖而出:

  1. 计算效率:
    自注意力的复杂度为 \( O(M^2 D) \) (与序列长度 \( M \) 二次方相关) ，而 Vim 的 SSM 为 \( O(M D N) \) (线性关系，其中 \( N \) 是一个较小固定值) 。
    自注意力机制与状态空间模型 (SSM) 的计算复杂度公式对比，突显二次方与线性扩展的区别。

  2. 内存效率:
    Vim 避免了存储庞大的注意力矩阵，使内存使用随序列长度线性增长。

  3. IO 效率:
    参数一次性加载到高速 GPU SRAM 中，计算在其中运行，最终输出才写回——最大限度减少数据传输瓶颈。

实验结果

论文在多种视觉任务上对 Vim 与领先模型进行了对比测试。

ImageNet 分类

一张表格，比较了各种卷积神经网络、Transformer 和基于 SSM 的模型在 ImageNet 上的 top-1 准确率。Vim 模型始终优于对应的 DeiT 模型。

表 1: Vim 始终优于同等规模的 DeiT 模型。长序列微调还能进一步提升准确率。

通过“长序列微调” (更小的图像块步幅) ，Vim 可以捕捉更细致的细节，准确率提升 1–2 个百分点。Vim-S† 达到 81.4%，与更大的 Transformer 模型相当。

高分辨率下的效率

折线图显示随着分辨率提升，Vim 的处理速度 (FPS) 始终高于 DeiT。

图 3: 低分辨率时速度相差不大，但分辨率越高，Vim 优势越明显。

折线图显示 GPU 内存使用量。DeiT 的显存占用随分辨率呈二次方增长并导致 OOM 错误，而 Vim 的显存使用呈线性增长。

图 4: 内存使用对比。Vim 在高分辨率下保持高效，避免 OOM 错误。

在 1248×1248 分辨率下，Vim 比 DeiT 快 2.8 倍，使用的 GPU 显存减少 86.8%，而 DeiT 则因 OOM (显存不足) 崩溃。

下游任务: 分割与检测

在 语义分割 (ADE20K) 上，Vim 优于 DeiT:

表格显示在 ADE20K 上的语义分割结果。Vim 在平均交并比 (mIoU) 上优于 DeiT。

表 2: Vim-S 的 mIoU 与 ResNet-101 相当，但参数量更少。

在 目标检测 和 实例分割 (COCO) 上，Vim-Ti 也大幅超越 DeiT-Ti:

表格显示在 COCO 上的目标检测和实例分割结果。Vim-Ti 的 AP 分数高于 DeiT-Ti。

表 3: 即使没有窗口化注意力等二维特定先验，Vim 的序列建模也优于 DeiT。

Vim 可直接完成高分辨率检测任务，无需结构调整，而 DeiT 则需要窗口化注意力。定性对比显示，Vim 能检测到 DeiT 遗漏的大型目标:

视觉对比图中，Vim-Ti 完整分割出一架大型 B-29 轰炸机，而 DeiT-Ti 只分割了部分。

图 5: 更强的长程上下文捕获能力让 Vim 能更完整地检测大型物体。

消融实验

双向性至关重要: 移除双向性会使分割性能显著下降 (-3.6 mIoU) 。采用完整的 双向 SSM + Conv1d 设计可在所有基准测试中取得最佳结果:

双向策略消融研究表格，完整设计在各项基准中取得最高准确率。

表 4: 双向处理对密集预测任务不可或缺。

[CLS] 令牌位置: 令人意外的是，将其放在序列中间效果最佳——准确率 76.1%——可充分利用双向处理优势:

分类设计消融研究表格，中间位置的分类令牌取得最高准确率。

表 5: 中间放置分类令牌可最大化前后向上下文的聚合。

结论: 新竞争者的出现

Vision Mamba 论文为基于 SSM 的视觉骨干提供了有力支持。通过引入双向处理和位置嵌入，Vim 实现了:

  1. 在多任务中匹敌或超越高度优化的 Transformer (如 DeiT) 。
  2. 速度和内存随序列长度线性扩展——特别适用于高分辨率输入。
  3. 保持纯序列建模，无需繁重的二维结构先验。

这种高效性让许多此前 Transformer 难以处理的任务成为可能: 如超高分辨率医学切片、超大卫星影像或长视频流——可端到端处理，无需复杂的分割拼接。

该论文表明，自注意力机制在高性能视觉中的统治地位并非不可撼动。有了 Vim，未来的计算机视觉骨干网络很可能将由状态空间模型来书写。