“Spatial-Mamba Effective Visual State Space Models via Structure-Aware State Fusion”

由香港理工大学、OPPO研究院、哈佛医学院和西安交通大学的研究团队联合提出，旨在解决现有视觉状态空间模型（SSMs）在处理2D图像时的局限性。论文提出Spatial-Mamba模型，通过结构感知状态融合在潜在状态空间中捕获相邻特征的空间依赖关系，解决现有视觉状态空间模型的局限，在多种视觉任务中取得优异表现。

论文地址：https://arxiv.org/abs/2410.15091

代码地址：https://github.com/EdwardChasel/Spatial-Mamba

研究动机

状态空间模型（SSMs）在自然语言处理中表现出色，Mamba等模型能有效捕捉1D序列长距离依赖。然而，视觉数据具有2D空间结构，现有视觉SSMs 常将图像转换为1D序列并采用扫描策略处理，存在诸多问题。如扫描会改变像素空间关系，破坏图像固有空间上下文；固定扫描路径难以捕捉复杂多变的空间关系，增加扫描方向又会导致计算量过大。因此，探索设计更有效、能感知结构的视觉SSMs迫在眉睫。

主要贡献

-提出新模型：提出Spatial-Mamba模型，引入结构感知状态融合（SASF）方程，直接在状态空间建立邻域连接，增强视觉上下文信息流动，有效捕捉图像空间结构依赖。

-统一理论框架：理论分析表明Spatial-Mamba、原始Mamba和线性注意力可统一在相同矩阵乘法框架下，加深了对模型工作机制的理解。

-性能优越：在图像分类、检测和分割等基础视觉任务上进行实验，结果显示Spatial-Mamba即使仅进行单次扫描，也能达到或超越基于不同扫描策略的最新模型性能。

方法

整体结构图：

-Spatial-Mamba公式化：在原始Mamba公式中引入SASF方程，整个过程由状态转移方程、SASF方程和观察方程描述。先通过单向扫描将输入图像转换为1D序列计算原始状态变量

，再用SASF方程融合相邻状态变量得到结构感知状态变量

，最后由观察方程生成输出

。实际中，利用多尺度扩张卷积对相邻状态变量线性加权，构建邻域集Ω。可视化结果表明，经SASF处理后的状态变量能更好地区分前景和背景，更准确地表示图像空间关系。

-网络架构：整体架构类似Swin-Transformer，包含四个连续阶段，有Spatial-Mamba-T、Spatial-Mamba-S和Spatial-Mamba-B三种不同尺度变体。输入图像经重叠茎层处理后得到特征图，再依次经过四个阶段，每个阶段包含多个Spatial-Mamba块和下采样层（最后阶段除外），最终由头部层处理特征得到任务输出。Spatial-Mamba块是基本构建单元，由结构感知SSM和前馈网络（FFN）组成，结构感知SSM通过将原始1D因果卷积替换为3×3深度卷积、将S6模块替换为SASF模块实现。

-与原始Mamba和线性注意力的联系：线性注意力、原始Mamba和Spatial-Mamba都可在统一矩阵乘法框架下建模，区别在于矩阵的结构。线性注意力中的亮值位置垂直方向一致，聚焦于小部分图像令牌；Mamba 的随时间呈衰减模式，能在先前图像令牌间转移焦点；Spatial-Mamba 的考虑更广泛空间邻域内所有状态的加权求和，能更全面地表示空间关系。

实验

-对比实验

 图像分类：Spatial-Mamba在ImageNet-1K分类任务中表现出色。Spatial-Mamba-T的Top-1准确率达到83.5%，超过ConvNeXt-T 1.4%，超越Swin-T 2.2%、NAT-T 0.3%，比VMamba-T和 LocalVMamba-T分别高出1.0%和0.8%。Spatial-Mamba-S和Spatial-Mamba-B的Top-1准确率分别为84.6%和85.3%，优于NAT-S、NAT-B、VMamba-S和VMamba-B。

 目标检测与实例分割：在COCO数据集的目标检测和实例分割任务里，不同调度下 Spatial-Mamba各变体模型均优于竞争对手。

 语义分割：在ADE20K数据集的语义分割任务中，Spatial-Mamba各变体模型表现优异。Spatial-Mamba-T单尺度交并比（mIoU）为48.6，多尺度mIoU为49.4，超过NAT-T和VMamba-T；Spatial-Mamba-B多尺度mIoU达到52.6，性能最佳。

-消融实验：对Spatial-Mamba-T进行消融实验，探究不同组件影响。调整邻域集Ω发现，从 3×3 邻域扩大到 5×5 邻域或采用扩张因子 d 为 1、3、5 的邻域集，虽会降低处理速度，但能提升精度。使用overlapped convolutions替换原始non-overlapped stem和下采样层，或融入局部感知单元（LPU），可丰富局部信息，提升模型精度。通过重新参数化技术实现结构感知状态融合（SASF）模块并优化 CUDA 内核，能加速模型；集成 MESA 可减轻过拟合，进一步提升模型性能。

仅供学习交流参考。

感谢阅读！可微信搜索公众号【AI启智汇】获取更多AI干货分享。