SF-Mamba：重新思考视觉 Mamba 的信息流与并行效率

Date 2026-07-09 · Category tech · Status finished · Confidence likely
论文解读, 视觉 Mamba, 状态空间模型

论文信息

论文题目：SF-Mamba: Rethinking State Space Model for Vision

论文地址：arXiv:2603.16423

这篇论文关注的不是“怎样把 Mamba 搬到视觉里”，而是一个更尖锐的问题：

现有视觉 Mamba 为什么理论上是线性的，实际却经常没有我们期待得那么快、那么强？

作者给出的答案有两部分：

  • 第一，单向 scan 有因果约束，前面的 patch 看不到后面的 patch。
  • 第二，视觉任务里的序列通常偏短，Mamba 的 scan kernel 在 GPU 上吃不满并行度。

因此，SF-Mamba 的核心不是再加更多大模块，而是直接瞄准这两个瓶颈：

  • 用 Auxiliary Patch Swapping 解决 future-to-past 信息流。
  • 用 Batch Folding with Periodic State Reset 解决短序列下的 GPU 利用率问题。

先看结论：它想优化的是速度-精度权衡

SF-Mamba 在 ImageNet-1K 上的精度-吞吐对比

这张图基本概括了整篇论文的目标。横轴是吞吐，纵轴是 Top-1 精度。作者追求的不是只提精度，也不是只冲速度，而是让模型整体往右上角移动。

和 MambaVision 相比，SF-Mamba 在三个模型规模上都同时变快、变准：

  • SF-Mamba-T: 82.5% / 7600 img/s，MambaVision-T: 82.3% / 6662 img/s
  • SF-Mamba-S: 83.5% / 5639 img/s，MambaVision-S: 83.3% / 4933 img/s
  • SF-Mamba-B: 84.4% / 3534 img/s，MambaVision-B: 84.2% / 2974 img/s

这也是全文最重要的主张：作者不是牺牲速度换精度，而是在更高吞吐下稳定拿到更好的结果。

动机：现有视觉 Mamba 的两个瓶颈

1. 单向扫描天然缺 future-to-past 信息流

Mamba 本质上是递推扫描。对图像 patch 序列来说，如果按从前到后的顺序扫，那么前面的 token 只能看到已经经过的位置，看不到后面区域的内容。

这在语言里很自然，在图像里却不自然。图像没有严格的因果顺序，左上角 patch 往往也需要知道右下角 patch 的上下文。

现有工作常用的补救方式是双向扫描或多方向扫描，但它们通常要做：

  • 双路甚至四路扫描
  • 反转或重排整段 token
  • 额外的数据搬运和实现开销

2. 视觉里的序列太短，GPU 并行吃不满

视觉 backbone 里，特别是深层 stage，patch 序列长度经常只有几十到几百。论文指出，Mamba 的 scan kernel 在这种短序列下效率并不理想，因为每条序列分到的 GPU 线程太少，线程利用率偏低。

这就是 SF-Mamba 的第二个切入口：不是 Mamba 理论不行，而是视觉任务的序列长度让工程实现没发挥出来。

第一个模块：Auxiliary Patch Swapping

Auxiliary Patch Swapping 示意图

这是论文最核心的建模模块，目的非常直接：

让单向 scan 也能把“后面的信息”传回“前面的 patch”。

它的做法并不复杂，但很巧。

它具体做了什么

给原始 patch 序列两端各加一个辅助 token：

  • x_head^aux
  • x_tail^aux

于是原来的序列

[x1, x2, ..., xT]

会变成

[x_head^aux, x1, x2, ..., xT, x_tail^aux]

在每个 stage 的第一个 Mamba block 里，这两个辅助 token 都用 avg(X) 初始化，也就是当前 patch 序列在序列维度上的平均特征。

为什么这能传递全局信息

因为 Mamba 还是按单向顺序扫描，最后那个 x_tail^aux 在这一层结束后，会得到一个新的输出 y_tail^aux。它在被更新时，已经“看过”整条序列，因此可以把它理解成当前层提炼出的全局摘要。

然后作者做了一步极轻量的 swap：

  • 下一层的头部辅助 token 用上一层的 y_tail^aux
  • 下一层的尾部辅助 token 用上一层的 y_head^aux

于是，上一层末端聚合出来的全局信息，在下一层会被送回最前面。这样前面的 patch 虽然在同一层里仍然是单向扫描，但在下一层一开始就能接触到“看过整张图”的全局提示。

这个设计为什么比双向扫描便宜

传统双向扫描的额外代价在于：

  • 序列要跑两遍
  • 还要对整段 token 做 O(n) 反转或重排

而这里：

  • 只多了两个辅助 token，长度从 T 变成 T + 2
  • 只交换两个 token，开销更接近 O(1)

所以它的本质可以概括为：

不直接做同层双向，而是把单向 scan 层末形成的全局摘要，跨层送回前端。

第二个模块：Batch Folding with Periodic State Reset

Batch Folding 与周期性状态重置

第二个模块不解决信息流，而是解决吞吐。

直觉很简单：把很多短序列临时拼成长序列

原始输入张量形状是：

[B, D, T]

作者把它 reshape 成：

[B1, D, (B2 * T)]，其中 B = B1 * B2

你可以把它理解成：把 batch 里的 B2 条短序列首尾拼成一条更长的“虚拟长序列”，让 scan kernel 更容易吃满 GPU。

但这样会带来信息泄漏

如果直接把不同样本拼起来递推，前一个样本的 hidden state 会流到下一个样本，导致不同图片之间串信息。

作者的解决办法是：

每隔 T 步就重置一次状态。

论文实现上是在边界处令状态转移里的 A_t = 0，等价于重新初始化 hidden state。这样虽然数据形式上被拼成了长序列，但每个样本在递推上仍然保持独立。

它为什么有效

论文在 SSM kernel 上直接测了加速收益：

不同 folding 比例下的 SSM speedup

作者报告，当把 batch 维“折叠”进 sequence 维后，SSM 计算部分能拿到大约 110% 到 180% 的 speedup，尤其在短序列设置下收益更明显。

还要额外处理卷积边界

Mamba block 里不只有 scan，还有 1D depthwise convolution。因此作者额外实现了适配 batch-folded 数据的卷积，确保卷积不会跨越不同样本的边界。

所以这个模块不是简单的 reshape trick，而是一套完整的执行优化：

  • 通过 folding 拉长有效序列
  • 通过 periodic reset 保持样本独立
  • 通过边界安全卷积维持等价计算

还有一个工程补充：Adaptive B1

Batch Folding 里并不是“拼得越多越好”。如果 folding 太弱，加速不明显；如果 folding 过猛，重排和边界处理的额外成本又会上来。

所以作者又加了一个 Adaptive B1 机制，根据 batch size、序列长度、模型维度和 state size，从离线构建的查找表里选一个更合适的 folding 比例。

这不是新的建模创新，但它很重要，因为它解释了为什么论文不仅在方法上成立，在真实吞吐上也能稳定跑出来。

实验 1：ImageNet-1K 分类

分类实验最能说明它是不是一个“普适 backbone 改进”，而不只是某个下游任务的小技巧。

论文里的关键数字如下：

  • SF-Mamba-T: 82.5, 7600 img/s
  • MambaVision-T: 82.3, 6662 img/s
  • SF-Mamba-S: 83.5, 5639 img/s
  • MambaVision-S: 83.3, 4933 img/s
  • SF-Mamba-B: 84.4, 3534 img/s
  • MambaVision-B: 84.2, 2974 img/s

这里最值得看的不是“精度只涨了 0.2”，而是：

  • 吞吐同步上升，且涨幅并不小
  • 这种趋势在 T/S/B 三个规模上都成立

如果再和更激进的多方向 Mamba 比较，像 Spatial-Mamba 的精度确实更高，但吞吐要低得多。比如：

  • Spatial-Mamba-T: 83.5, 1430 img/s
  • Spatial-Mamba-B: 85.3, 670 img/s

这也再次印证了论文的主题：作者主要优化的是速度-精度权衡，而不是孤立地追求最高精度。

实验 2：ADE20K 语义分割

语义分割更依赖像素级边界和全局结构理解，所以很适合检验 future-to-past 信息流是否真的有用。

论文在 UPerNet 上的主结果如下：

  • Swin-T: 44.5 mIoU, 40.0 fps

  • Focal-T: 45.8 mIoU, 38.9 fps

  • MambaVision-T: 46.0 mIoU, 45.0 fps

  • SF-Mamba-T: 47.2 mIoU, 47.9 fps

  • MambaVision-S: 48.2 mIoU, 40.9 fps

  • SF-Mamba-S: 48.5 mIoU, 45.4 fps

  • MambaVision-B: 49.1 mIoU, 37.3 fps

  • SF-Mamba-B: 50.1 mIoU, 42.6 fps

这组结果很有说服力，因为它说明在更依赖空间结构的任务上，辅助 token swap 带来的全局上下文回流确实能转化成下游收益。

论文还给了一个使用 windowed attention 的轻量版本，文中记作 SF-Mamba♣。它的特点是：

  • Tiny: 46.5 mIoU, 48.7 fps
  • Small: 48.1 mIoU, 47.3 fps
  • Base: 49.1 mIoU, 42.7 fps

也就是说，如果更关心高分辨率下的推理代价，还可以进一步拿窗口化 attention 压 FLOPs，而不会明显破坏整体权衡。

实验 3：COCO 检测与实例分割

Cascade Mask R-CNN，3x schedule

论文在 COCO 2017 上给出了更完整的对比。

Tiny 级别：

  • ConvNeXt-T: 50.4 AP^b / 43.7 AP^m, 32.1 fps
  • MambaVision-T: 51.1 / 44.3, 19.4 fps
  • SF-Mamba-T: 51.0 / 44.2, 27.8 fps

Small 级别：

  • ConvNeXt-S: 51.9 / 45.0, 28.0 fps
  • MambaVision-S: 52.3 / 45.2, 20.5 fps
  • SF-Mamba-S: 52.4 / 45.4, 28.7 fps

Base 级别：

  • ConvNeXt-B: 52.7 / 45.6, 26.0 fps
  • MambaVision-B: 52.8 / 45.7, 16.4 fps
  • SF-Mamba-B: 52.8 / 45.8, 26.8 fps

这组结果尤其能说明 SF-Mamba 的工程价值：它基本把 MambaVision 在 COCO 上过高的计算代价拉回到了更合理的位置，同时保持甚至略微提升精度。

Faster R-CNN，1x schedule

论文还补了另一个检测头，结果同样一致：

  • MambaVision-T: 42.4 AP, 37.2 fps
  • SF-Mamba-T: 43.2 AP, 41.7 fps
  • MambaVision-S: 43.9 AP, 37.2 fps
  • SF-Mamba-S: 44.9 AP, 40.8 fps
  • MambaVision-B: 46.2 AP, 30.0 fps
  • SF-Mamba-B: 47.6 AP, 34.8 fps

Mask R-CNN，1x schedule

这里作者还拿它和 VMamba 做了性能-吞吐对比：

  • VMamba-T: 47.3 AP^b / 42.7 AP^m, 29.9 fps

  • SF-Mamba-T: 43.8 / 40.3, 40.9 fps

  • VMamba-S: 48.7 / 43.7, 23.3 fps

  • SF-Mamba-S: 45.3 / 41.5, 41.0 fps

  • VMamba-B: 49.2 / 44.1, 20.0 fps

  • SF-Mamba-B: 47.8 / 43.5, 34.6 fps

它没有在同规模上全面压过 VMamba 的绝对精度，但速度优势很明显，所以作者强调的是 performance-throughput trade-off，不是单点精度冠军。

补充实验最关键的一点：swap 本身就能带来稳定收益

附录里的一个结果我觉得特别重要，因为它直接回答了“增益到底来自哪个模块”。

在相同 FLOPs 下：

  • w/o swap: 46.0 mIoU, 50.9 mAP^b, 44.1 mAP^m
  • w/ swap: 47.2 mIoU, 51.2 mAP^b, 44.5 mAP^m

也就是说，哪怕先不看 batch folding，光是辅助 token swapping 本身，就已经在分割和检测上稳定把性能往上推了一步。

这说明第一模块不是“锦上添花”，而是真正解决了视觉 Mamba 的因果约束问题。

我对这篇论文的判断

我觉得这篇工作的优点非常明确：

  • 它抓的问题很准，不是泛泛地说“再做一个更强的视觉 Mamba”
  • 两个核心模块都很有针对性，一个改信息流，一个改执行效率
  • 主结果和补充实验是一致的，分类、分割、检测都支持同一条叙事

它的边界也同样明确：

  • 它不是要拿绝对最高精度，而是要拿更好的速度-精度权衡
  • 它大量收益依赖于具体工程实现，尤其是 kernel 和卷积边界处理
  • 它并没有完全替代所有多方向或高精度视觉 Mamba，只是在“高效 backbone”这条线上很强

如果用一句话概括，我会这么评价：

SF-Mamba 不是重新发明视觉 Mamba，而是把现有视觉 Mamba 最不舒服的两个地方，分别用一个很便宜、很实用的模块补上了。

总结

SF-Mamba 最值得记住的，不是它提出了多复杂的数学，而是它把视觉 Mamba 的两个真实瓶颈讲清楚了：

  • 前面的 patch 看不到后面的 patch
  • 短序列下 GPU 吃不满 scan 并行

对应地，作者给出的答案也很清晰：

  • 用 Auxiliary Patch Swapping 把层末的全局摘要跨层送回前端
  • 用 Batch Folding with Periodic State Reset 把很多短序列临时拼成长序列，再通过状态重置保证样本独立

最后的实验结果也支持这条逻辑：在 ImageNet-1K、ADE20K 和 COCO 上，SF-Mamba 基本都能在保持甚至提升精度的同时，拿到更高吞吐。

如果你现在正关注 Vision Mamba / Vision SSM 这条路线，这篇论文很值得看，因为它提供的不是“再堆一个模块”的启发，而是：

视觉 SSM 的能力，不只取决于状态空间模型本身，还取决于你怎样组织信息流，以及怎样把线性复杂度真正落到 GPU 上。

See also