SF-Mamba：高效视觉Mamba与辅助Patch交换

引言

视觉Transformer（ViT）尽管在各类视觉任务中取得了显著成功，但其自注意力机制的计算复杂度随序列长度呈二次方增长，这给高分辨率图像处理带来了沉重的计算负担。Mamba通过选择性状态空间模型（Selective State Space Model, SSM）实现了线性时间复杂度的序列建模，在视觉领域展现出巨大潜力。

然而，现有视觉Mamba方法面临两个关键瓶颈：

因果性约束：Mamba的单向扫描机制限制了patch间的双向信息交互
GPU并行效率低：短序列场景下GPU线程利用率不足，导致推理速度慢

SF-Mamba（Spatial-Frequency Mamba）由Sony集团的研究者提出，从数据流和计算效率两个维度重新思考视觉SSM设计，是对视觉Mamba模型的重要改进，实现了优于现有方法的精度-吞吐量权衡。¹

核心创新

SF-Mamba提出两项核心技术：

创新点	目标	方法
辅助Patch交换	双向信息流	引入两个辅助token实现未来→过去的信息路由
批量折叠	GPU并行优化	合并batch和sequence维度，周期性重置状态

辅助Patch交换机制

问题背景

图像patch之间不存在严格的因果顺序，限制Mamba块为单向扫描会阻碍信息流动：

┌─────────────────────────────────────────────────────────────┐
│ 单向扫描的信息限制                                           │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│   图像:                    单向扫描顺序:                    │
│   ┌───┬───┬───┐          1 → 2 → 3 → 4 → 5 → 6 → 7 → 8   │
│   │ 1 │ 2 │ 3 │ 4                                               │
│   ├───┼───┼───┤          问题: 早期patch无法直接访问           │
│   │ 5 │ 6 │ 7 │ 8        后期patch的信息                       │
│   └───┴───┴───┘                                              │
│                                                             │
│   例如: patch 1 无法直接获取 patch 8 的信息                    │
│                                                             │
└─────────────────────────────────────────────────────────────┘

现有方法通过多方向扫描（如双向扫描、十字扫描）解决此问题，但存在显著开销：

数据重排开销：token重排占推理时间的5-8%
并行扫描开销：多方向并行处理占推理时间的28-42%
2D/1D格式切换：VMamba等方法需在二维和一维格式间频繁转换

方法原理

SF-Mamba提出**辅助patch交换（Auxiliary Patch Swapping）**机制，仅用两个额外token实现双向信息流：

┌─────────────────────────────────────────────────────────────┐
│ 辅助Patch交换机制                                           │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│   输入序列:    [x_aux_head, x₁, x₂, ..., xₜ, x_aux_tail]   │
│                                                             │
│   单向扫描处理:                                             │
│   [x_aux_head] → [x₁] → [x₂] → ... → [xₜ] → [x_aux_tail] │
│        ↓              ↓        ↓             ↓             │
│       h₀            h₁       h₂            hₜ₊₁           │
│                                                             │
│   核心洞察:                                                  │
│   - x_aux_tail 能从所有patch中提取全局信息                   │
│   - 通过交换，x_aux_tail 的信息传递给下一个block的head       │
│                                                             │
│   交换操作（第i层到第i+1层）:                               │
│   x_aux,i+1^head = y_aux,i^tail                              │
│   x_aux,i+1^tail = y_aux,i^head                              │
│                                                             │
└─────────────────────────────────────────────────────────────┘

数学推导：

设输入序列为 $X = (x_{1}, x_{2}, ..., x_{T})$ ，在第一层Mamba block前添加两个辅助token：

X^{'} = (x_{aux, 1}^{head}, x_{1}, x_{2}, ..., x_{T}, x_{aux, 1}^{tail})

经过第 $i$ 层Mamba块处理后，两个辅助token的输出分别为 $y_{aux, i}^{head}$ 和 $y_{aux, i}^{tail}$ ，其中：

$y_{aux, i}^{tail}$ 包含了从 $t = 0$ 到 $t = T$ 的全局信息
通过交换，下一层的patch序列可以”看到”全局信息

交换公式为：

x_{aux, i + 1}^{head} = y_{aux, i}^{tail}, x_{aux, i + 1}^{tail} = y_{aux, i}^{head}

与多方向扫描的对比

扫描方式	token交换复杂度	额外计算	精度提升
双向扫描	$O (n)$ （全序列翻转）	2× SSM计算	高
十字扫描	$O (n)$ （四方向重排）	4× SSM计算	高
辅助交换	$O (1)$	2个token	相当

有效感受野分析

通过有效感受野（Effective Receptive Field, ERF）分析验证双向信息流的有效性：

┌─────────────────────────────────────────────────────────────┐
│ 有效感受野对比                                              │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│   MambaVision (单向扫描):                                   │
│   - 仅能通过卷积获取局部上下文                               │
│   - patch 1 对图像下半部分响应弱                            │
│                                                             │
│   SF-Mamba (辅助交换):                                      │
│   - 通过辅助token实现全局信息传递                            │
│   - patch 1 对整幅图像响应均匀                              │
│                                                             │
│   ERF可视化显示: SF-Mamba在浅层即实现全局感受野              │
│                                                             │
└─────────────────────────────────────────────────────────────┘

批量折叠与周期性状态重置

问题背景

Mamba采用CUDA warp-scan函数实现并行扫描，该函数以32线程为一组进行处理。在视觉任务中，序列长度通常较短（如MambaVision Stage 3为196，Stage 4为49），导致GPU线程利用率低下。

┌─────────────────────────────────────────────────────────────┐
│ 序列长度与线程利用率                                        │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│   GPU warp (32 threads):                                    │
│   ┌─────────────────────────────────────────────────────┐ │
│   │ [T0][T1][T2]...[T30][T31]                          │ │
│   └─────────────────────────────────────────────────────┘ │
│        ↓                                                   │
│   序列长度T=196时: 32线程/warp，利用率 = 196/32 = 6.1 warp │
│   序列长度T=49时:  32线程/warp，利用率 = 49/32  = 1.5 warp │
│                                                             │
│   问题: 大量线程处于空闲状态                                 │
│                                                             │
└─────────────────────────────────────────────────────────────┘

方法原理

SF-Mamba提出**批量折叠（Batch Folding）**策略，将batch维度和sequence维度合并：

设原始输入张量为 $Z \in R^{B \times D \times T}$ ，其中 $B$ 为batch大小， $D$ 为维度， $T$ 为序列长度。折叠操作定义为：

Z^{'} \in R^{B_{1} \times D \times (B_{2} \cdot T)}, B = B_{1} \cdot B_{2}

┌─────────────────────────────────────────────────────────────┐
│ 批量折叠操作                                                │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│   原始:                                                     │
│   Z ∈ ℝ^(B × D × T)                                        │
│   ┌────────────────────────────────────────────────────┐   │
│   │ seq 1: [t₀, t₁, ..., tₜ]                          │   │
│   │ seq 2: [t₀, t₁, ..., tₜ]                          │   │
│   │ ...                                                │   │
│   │ seq B: [t₀, t₁, ..., tₜ]                          │   │
│   └────────────────────────────────────────────────────┘   │
│                                                             │
│   折叠后:                                                   │
│   Z' ∈ ℝ^(B₁ × D × (B₂·T))                               │
│   ┌────────────────────────────────────────────────────┐   │
│   │ [seq1_t₀, seq1_t₁, ..., seq1_tₜ, seq2_t₀, ...]    │   │
│   └────────────────────────────────────────────────────┘   │
│   ↑                                                        │
│   B₂个短序列拼接成一个长序列                                │
│                                                             │
└─────────────────────────────────────────────────────────────┘

周期性状态重置

折叠操作会混合不同序列的隐藏状态，导致信息泄露。为解决此问题，SF-Mamba采用周期性状态重置（Periodic State Reset）：

h_{t} \leftarrow A_{t} h_{t - 1} + B_{t} x_{t}

当 $t mod T = 0$ 时，设置 $A_{t} = 0$ ，这相当于重置隐藏状态：

h_{t} \leftarrow B_{t} x_{t} (当 t mod T = 0 时)

由于 $B_{t}$ 和 $C_{t}$ 仅作用于当前输入和隐藏状态，无需重置。该操作仅需重置 $A_{t}$ ，计算开销极小。

┌─────────────────────────────────────────────────────────────┐
│ 周期性状态重置                                              │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│   折叠序列中的处理:                                          │
│   [seq₁_t₀][seq₁_t₁]...[seq₁_tₜ]|[seq₂_t₀][seq₂_t₁]...    │
│   ↑         ↑              ↑          ↑                    │
│   A=0      A≠0           A=0       A=0  ← 状态重置点       │
│   (reset)                        (reset)                   │
│                                                             │
│   结果: 折叠后的输出等价于未折叠的输出                       │
│                                                             │
└─────────────────────────────────────────────────────────────┘

自适应B₁选择

折叠比例 $B_{1} / B$ 的最优值由多种因素决定（batch大小、序列长度、维度、状态维度等）。SF-Mamba采用预计算的4维查找表（LUT）：

B_{1} = f (B, B \cdot LUT (B, D, S, L))

其中 $f (a, b)$ 返回 $a$ 的约数中最接近 $b$ 的值。

1D深度可分离卷积适配

批量折叠后，需要适配Mamba块中的1D深度可分离卷积。实现要点：

在每段 $T$ 序列边界处进行隐式填充
确保卷积不会跨越不同序列的边界
使用Triton CUDA kernel优化交换操作，约提升40 img/s（tiny模型）

架构设计

SF-Mamba采用MambaVision作为宏观架构基础，结合Mamba和Transformer的混合设计。

四阶段层级结构

┌─────────────────────────────────────────────────────────────┐
│ SF-Mamba 宏观架构                                           │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│   输入图像 I ∈ ℝ^(H×W×3)                                    │
│       ↓                                                     │
│   ┌─────────────────────────────────────────────────────┐   │
│   │ Stage 1: Conv Stem + ConvBlock × N₁                │   │
│   │ 输出: H/4 × W/4 × D                                 │   │
│   └─────────────────────────────────────────────────────┘   │
│       ↓                                                     │
│   ┌─────────────────────────────────────────────────────┐   │
│   │ Stage 2: Downsample + ConvBlock × N₂               │   │
│   │ 输出: H/8 × W/8 × 2D                                │   │
│   └─────────────────────────────────────────────────────┘   │
│       ↓                                                     │
│   ┌─────────────────────────────────────────────────────┐   │
│   │ Stage 3: Downsample + MambaBlock×(N₃/2) +         │   │
│   │           AttenBlock×(N₃/2)                         │   │
│   │ 输出: H/16 × W/16 × 4D  [含辅助token交换]          │   │
│   └─────────────────────────────────────────────────────┘   │
│       ↓                                                     │
│   ┌─────────────────────────────────────────────────────┐   │
│   │ Stage 4: Downsample + MambaBlock×(N₄/2) +         │   │
│   │           AttenBlock×(N₄/2)                         │   │
│   │ 输出: H/32 × W/32 × 8D  [含辅助token交换]          │   │
│   └─────────────────────────────────────────────────────┘   │
│       ↓                                                     │
│   Global AvgPool + Linear → 类别预测                        │
│                                                             │
└─────────────────────────────────────────────────────────────┘

MambaVision Mixer

SF-Mamba的Mamba块采用MambaVision Mixer结构，包含两个并行分支：

X_{1} = SSM (σ (Conv (Linear_{D \to D} (X_{in})))) X_{2} = σ (Conv (Linear_{D \to D} (X_{in}))) Y = Linear_{2 D \to D} (Concat (X_{1}, X_{2}))

其中 $σ$ 为SiLU激活函数，Conv为1D深度可分离卷积。

辅助token初始化与丢弃

方面	设计选择
初始化	数据依赖型（全局平均池化），优于可学习型
丢弃时机	第一个Attention层之后
计算开销	仅增加2个token，几乎可忽略

实验结果

图像分类（ImageNet-1K）

模型	参数量	GFLOPs	吞吐量 (img/s)	Top-1 Acc
Swin-T	29M	4.5G	2863	81.3%
VMamba-T	30M	4.9G	1684	82.6%
Spatial-Mamba-T	27M	4.5G	1430	83.5%
MambaVision-T	32M	4.4G	6662	82.3%
SF-Mamba-T	32M	4.5G	7600	82.5%
MambaVision-S	50M	7.5G	4933	83.3%
SF-Mamba-S	50M	7.6G	5639	83.5%
SF-Mamba-B	98M	15.1G	3534	84.4%

关键发现：

SF-Mamba在所有规模下均实现更高的吞吐量
精度与MambaVision相当甚至更好
在精度-吞吐量权衡曲线上处于Pareto前沿

语义分割（ADE20K + UperNet）

模型	参数量	mIoU	FPS
Swin-T	60M	44.5%	40.0
MambaVision-T	62M	46.0%	45.0
SF-Mamba-T	62M	47.2%	47.9
MambaVision-B	130M	49.1%	37.3
SF-Mamba-B	130M	50.1%	42.6

目标检测与实例分割（COCO + Cascade Mask R-CNN）

模型	box AP	mask AP	FPS
Swin-T	50.4%	47.3%	-
MambaVision-T	51.1%	47.9%	19.4
SF-Mamba-T	51.0%	48.0%	27.8
SF-Mamba♣-T	50.9%	47.7%	28.3

注：SF-Mamba♣ 使用窗口注意力以降低计算量

消融实验

辅助token交换的效果

配置	IN1K Acc	ADE20K mIoU	吞吐量
无交换，仅添加可学习token	82.1%	46.2%	7613
交换+数据依赖初始化	82.5%	47.2%	7600
无交换	82.2%	46.0%	7645

结论：性能提升来自双向信息流而非额外token的灵活性。

批量折叠的效果

配置	吞吐量提升
MambaVision-T基线	-
+实现优化	+5%
+批量折叠	+14%
+自适应B₁	+2%
+1D卷积适配	+2%
总计	+14%

与相关工作的对比

多方向扫描的开销分析

方法	token重排开销	并行路径开销
Vim-S	5.4%	28.0%
VMamba-T	7.8%	42.3%
SF-Mamba-T	<1%	-

与其他视觉Mamba变体的对比

变体	特点	局限性
Vim	双向SSM扫描	需要2×计算量
VMamba	十字4方向扫描	2D/1D格式切换开销
MambaVision	Mamba+Attention混合	浅层缺乏未来信息
SF-Mamba	辅助token交换+批量折叠	兼容性问题（含2D卷积的架构）

技术细节

实现优化

移除未使用的row-dimension chunking：减少中间变量的内存访问
推理时抑制隐藏状态输出：减少不必要的内存写入
用1D逐点卷积替代线性层：减少张量重排
Triton CUDA kernel优化：加速token交换操作

适用性分析

SF-Mamba的两项技术可应用于：

✅ MambaVision架构：直接适用
✅ Vim架构：性能大幅提升
❌ 含2D卷积的架构（如VMamba）：需要每层恢复2D格式，速度提升有限

在Vim-S上的应用效果：

配置	参数量	GFLOPs	吞吐量	Top-1 Acc
Vim-S (双向)	26M	5.3G	1079	80.3%
Vim-S (单向)	26M	4.9G	1639	79.3%
+交换	26M	5.0G	1614	80.1%
+交换+折叠	26M	5.0G	2022	80.1%

结论

SF-Mamba从两个新视角重新思考视觉状态空间模型：

数据流视角：辅助patch交换机制以 $O (1)$ 的token交换开销实现双向信息流，弥补了单向扫描的局限性
计算效率视角：批量折叠通过合并batch和sequence维度最大化GPU利用率，周期性状态重置确保计算正确性

在ImageNet分类、COCO目标检测和ADE20K语义分割任务上，SF-Mamba均实现了精度和吞吐量的双重提升，为构建真正高效的视觉Mamba骨干网络提供了新思路。

SF-Mamba: Rethinking State Space Model for Vision. Masakazu Yoshimura, Teruaki Hayashi, Yuki Hoshino, Wei-Yao Wang, Takeshi Ohashi. Sony Group Corporation, 2026. ↩

Metaphor

探索

SF-Mamba：高效视觉Mamba与辅助Patch交换

引言

核心创新

辅助Patch交换机制

问题背景

方法原理

与多方向扫描的对比

有效感受野分析

批量折叠与周期性状态重置

问题背景

方法原理

周期性状态重置

自适应B₁选择

1D深度可分离卷积适配

架构设计

四阶段层级结构

MambaVision Mixer

辅助token初始化与丢弃

实验结果

图像分类（ImageNet-1K）

语义分割（ADE20K + UperNet）

目标检测与实例分割（COCO + Cascade Mask R-CNN）

消融实验

辅助token交换的效果

批量折叠的效果

与相关工作的对比

多方向扫描的开销分析

与其他视觉Mamba变体的对比

技术细节

实现优化

适用性分析

结论

关系图谱

目录

Metaphor

探索

SF-Mamba：高效视觉Mamba与辅助Patch交换

引言

核心创新

辅助Patch交换机制

问题背景

方法原理

与多方向扫描的对比

有效感受野分析

批量折叠与周期性状态重置

问题背景

方法原理

周期性状态重置

自适应B₁选择

1D深度可分离卷积适配

架构设计

四阶段层级结构

MambaVision Mixer

辅助token初始化与丢弃

实验结果

图像分类（ImageNet-1K）

语义分割（ADE20K + UperNet）

目标检测与实例分割（COCO + Cascade Mask R-CNN）

消融实验

辅助token交换的效果

批量折叠的效果

与相关工作的对比

多方向扫描的开销分析

与其他视觉Mamba变体的对比

技术细节

实现优化

适用性分析

结论

Footnotes

关系图谱

目录