多年来，计算机视觉领域一直由两大架构巨头主导:** 卷积神经网络 (CNNs)** 和 视觉 Transformer (ViTs)。CNN 通过滑动卷积核擅长捕捉局部特征，而 ViT 则利用自注意力机制来建模整个图像的全局关系。

现在，一个新的竞争者从序列建模的世界中脱颖而出:** 状态空间模型 (SSM)**，特别是其现代高性能变体 Mamba。

Mamba 在处理如文本和基因组学等长一维序列方面展现了非凡的能力，提供了线性时间复杂度和令人印象深刻的性能。研究人员自然希望将其优势引入视觉任务中。然而，尽管 Vision Mamba (Vim) 和 VMamba 等早期尝试很有前景，但它们并未决定性地超越 CNN 和 ViT。这就提出了一个关键问题:

为什么一个在 1D 数据上如此强大的模型，在 2D 图像上却难以发挥其全部潜力？

最近的一篇论文《LocalMamba: Visual State Space Model with Windowed Selective Scan》指出，答案不在于模型本身，而在于我们如何向其输入视觉信息。核心挑战在于一种不匹配: 图像是具有强局部相关性的二维像素网格，而 Mamba 是一个一维序列处理器。将图像展平为一个长序列会破坏这些至关重要的局部关系。

LocalMamba 引入了一个巧妙的修正方法:

  1. 窗口化选择性扫描，保留局部二维依赖关系。
  2. 为每一层自动搜索最佳扫描模式。

这不仅仅是一项渐进式的改进——它从根本上重新思考了 SSM 应该如何看待图像，从而在分类、检测和分割任务中取得了显著的进步。

路线图

在本文中，我们将探讨:

  • 状态空间模型的基础知识及其效率。
  • 将 1D 序列模型应用于 2D 图像的核心问题。
  • LocalMamba 的优雅解决方案: 局部扫描 + 自动化方向搜索。
  • 证明该方法是视觉 SSM 迈出的重要一步的实验结果。

背景: 从卷积到状态空间

状态空间模型 (SSM) 简介

状态空间模型源于控制理论，后被应用于深度学习中进行序列建模。它们通过一个中间的潜在状态来描述序列的演变:

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

\[ y(t) = \boldsymbol{C}h(t) \]

这里:

  • \(x(t)\): t 时刻的输入
  • \(y(t)\): t 时刻的输出
  • \(h(t)\): 隐藏状态
  • \(\boldsymbol{A}, \boldsymbol{B}, \boldsymbol{C}\): 定义系统动态的矩阵

离散化后:

\[ h_t = \bar{A} h_{t-1} + \bar{B} x_t \]

\[ y_t = C h_t \]

这类似于一个 RNN，但经典的 SSM 可以通过全局卷积重构以进行并行计算:

\[ y = x \circledast \overline{K} \]

其中:

\[ \overline{K} = (C\overline{B}, C\overline{AB}, \dots, C \overline{A}^{L-1}\overline{B}) \]

这使得模型能够像 CNN 一样实现线性时间复杂度的扩展和高效训练。

Mamba: 让 SSM 具有选择性

传统 SSM 使用静态参数 (\(\bar{A}, \bar{B}, C\))，不随输入变化。Mamba 引入了依赖于输入的参数化 (“选择性扫描”或 S6) ，从序列本身生成 \(\boldsymbol{B}, \boldsymbol{C}\) 和一个时间尺度参数 \(\Delta\)。

这种动态适应性使 Mamba 能够根据内容选择性地存储或遗忘信息，从而在语言建模中取得了顶尖成果。

核心问题: 方枘圆凿

Mamba 处理的是序列。但图像是二维网格——我们如何将它们输入到一个序列模型中呢？

最直接的方法是:** 将图像展平为补丁 (patches)**，并排列成一个单一的序列。

不同扫描方法的图示。左侧显示了覆盖在一条蛇图像上的网格。面板 (a) 显示逐行扫描。面板 (b) 显示逐列扫描。面板 (c) 显示 LocalMamba 的窗口化扫描，保持了邻近像素的紧密性。

图 1: 像 Vim (a) 这样的先前方法逐行展开图像，会使垂直相邻的像素被隔得很远。VMamba (b) 使用交叉扫描，但仍会在同一次扫描中处理相距较远的区域。LocalMamba (c) 将图像划分为窗口并在窗口内扫描，从而保留局部关系。

展平操作破坏了空间局部性: 相邻像素在 1D 序列中可能相隔数百步。这迫使 SSM 去学习那些 CNN 能够轻松检测到的局部模式。

VMamba 通过水平和垂直交叉扫描有所改善。但即使如此，它仍会将相距很远的标记 (token) 一起处理，难以完全保留局部细节。

LocalMamba 方法: 在窗口中观察

LocalMamba 通过以下方式解决了这个问题:

  1. 局部扫描 (Local Scanning): 保留邻域关系
  2. 空间与通道注意力 (SCAttn): 智能融合特征
  3. 自动化扫描方向搜索: 为每一层找到最优扫描方式

1. 局部扫描: 保留邻域

LocalMamba 不是一次性扫描整幅图像，而是将其划分为**不重叠的窗口 **(如 \(3 \times 3\) 或 \(7 \times 7\) 补丁) ，并分别扫描每个窗口。

原图中相近的标记在处理后的序列中仍保持接近，从而高效捕捉局部细节。

为保留全局上下文，LocalMamba 模块使用四个并行扫描:

  • 标准全局扫描
  • 局部窗口扫描
  • 上述两种扫描的翻转版本 (逆序)

LocalVim 模块和 SCAttn 模块的架构。面板 (a) 展示了四个并行扫描分支，其输出送入 SSM 和 SCAttn 模块。面板 (b) 展示了 SCAttn 模块，它使用通道和空间注意力为特征加权。

图 2: (a) LocalVim 模块: 四个扫描分支并行处理。(b) 空间与通道注意力 (SCAttn) 模块融合这些特征，通过对通道和空间位置加权实现最优融合。

2. SCAttn: 智能融合扫描结果

多个扫描会生成不同特征，关键在于如何在不丢失重要信息的前提下组合它们。

SCAttn 通过学习:

  • 通道注意力: 识别信息量最大的特征通道
  • 空间注意力: 识别最显著的标记 (补丁)

确保局部和全局信息融合成统一且丰富的表示。

3. 自动化扫描方向搜索

每层最优的扫描方向——水平、垂直、局部小窗口、局部大窗口——可能不同。

手工设计既耗时又难以最优。受 DARTS 启发，LocalMamba 对 8 个候选方向进行可微搜索:

  • 水平
  • 垂直
  • \(2 \times 2\) 局部
  • \(7 \times 7\) 局部
  • 以上各方向的翻转版本

每层学习所有候选方向的加权和:

\[ \boldsymbol{y}^{(l)} = \sum_{s \in \mathcal{S}} \frac{\exp(\alpha_s^{(l)})}{\sum_{s' \in \mathcal{S}} \exp(\alpha_{s'}^{(l)})} \mathrm{SSM}_s(\boldsymbol{x}^{(l)}) \]

在训练完“超网”后，每层保留得分最高的 4 个方向用于最终模型。

由自动搜索选择的不同模型扫描方向可视化。

图 3: LocalVim 的搜索模式。局部扫描 (橙色、红色) 在早期和后期层占主导，全局扫描 (蓝色、绿色) 更多出现在中间层。

实验: 检验 LocalMamba

两个系列:

  • LocalVim: 采用与原始 Vim 相同的均匀结构
  • LocalVMamba: 采用类似 Swin Transformer 的分层结构

ImageNet 分类

在 ImageNet-1K 分类准确率上比较不同骨干模型的表格。

表 1: ImageNet-1K 分类结果。* 表示使用局部扫描但未进行自动搜索的模型。

LocalVim-T 在 FLOPs 相同 (1.5G) 的情况下，准确率达到 76.2%，比 Vim-Ti 提升 +3.1%。即使是未搜索的变体 LocalVim-T* 也达 75.8%，表明局部扫描本身贡献了大部分改进。

条形图显示了向 Vim 模型添加局部扫描带来的准确率提升。

图 4: 向 Vim 添加局部扫描显著提升准确率。

COCO 目标检测

COCO 物体检测和分割结果表。

表 2: 使用 Mask R-CNN 在 COCO 上的结果。LocalVMamba 击败了 VMamba 及其他强劲骨干网络，如 Swin。

LocalVMamba-T 在 box AP 上比 Swin-T 提升 +4.0，同时在 mask AP 上有类似增益。

ADE20K 语义分割

ADE20K 分割结果表。

表 3: ADE20K 分割结果。LocalVim-S 比 Vim-S 的 mIoU 高 1.5。

LocalVMamba-S 达到 51.0 mIoU (MS)，创下基于 Mamba 的视觉骨干网络新基准。

消融实验

消融研究表，显示了每个组件的贡献。

表 4: 局部扫描、组合扫描、SCAttn 每个组件都对准确率提升有贡献。

逐步分析:

  1. 基线 (Vim-T): 73.1%
  2. 仅局部扫描: +1.0%
  3. 局部 + 全局扫描: +1.1%
  4. 加入 SCAttn: +0.6%
  5. 加入搜索: +0.4%

结论

LocalMamba 通过解决 1D 序列处理与 2D 图像结构的不匹配问题，为使 SSM 在视觉领域有效提供了有力答案。

主要贡献:

  1. 窗口化局部扫描: 保留空间邻域，促进精细特征学习
  2. 自动化方向搜索: 为每层找到最优的局部/全局扫描组合

结果清晰: 相较之前的 Vim 和 VMamba 模型有显著提升，可与顶尖 CNN 和 ViT 竞争，并在部分场景下取得了新的 SOTA 水平。

对于视觉 SSM 而言，如何扫描与扫描什么同等重要。

通过给予序列模型更自然的“看”图像方式，LocalMamba 为将其他领域的强大架构应用于视觉挑战开辟了令人振奋的新方向——在效率与精度之间取得平衡，或将重塑未来深度学习骨干网络。