Spatial-Mamba: Effective Visual State Space Models via Structure-aware State Fusion
作者: Chaodong Xiao, Minghan Li, Zhengqiang Zhang, Deyu Meng, Lei Zhang
分类: cs.CV
发布日期: 2024-10-19 (更新: 2025-02-26)
备注: Accepted by ICLR 2025
🔗 代码/项目: GITHUB
💡 一句话要点
提出Spatial-Mamba，通过结构感知状态融合有效建模视觉状态空间，提升图像理解能力。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 视觉状态空间模型 空间依赖建模 扩张卷积 图像分类 目标检测 语义分割 Mamba 结构感知
📋 核心要点
- 现有视觉SSM通过扫描图像序列建模空间关系，但难以有效捕获复杂空间结构，且计算成本高昂。
- Spatial-Mamba直接在状态空间建立邻域连接，利用扩张卷积进行结构感知状态融合，增强上下文信息流动。
- 实验表明，Spatial-Mamba在图像分类、检测和分割任务上，单次扫描即可达到或超越现有SOTA模型。
📝 摘要（中文）
选择性状态空间模型(SSMs)，如Mamba，在捕获一维序列数据的长程依赖方面表现出色，但将其应用于二维视觉任务仍然面临挑战。现有的视觉SSM通常将图像转换为一维序列，并采用各种扫描模式来结合局部空间依赖性。然而，这些方法在有效捕获复杂的图像空间结构方面存在局限性，并且由于扫描路径的延长而导致计算成本增加。为了解决这些限制，我们提出Spatial-Mamba，一种新颖的方法，可以直接在状态空间中建立邻域连接。我们没有仅仅依赖于顺序状态转换，而是引入了一种结构感知的状态融合方程，该方程利用扩张卷积来捕获图像空间结构依赖性，从而显著增强了视觉上下文信息的流动。Spatial-Mamba分三个阶段进行：单向扫描中的初始状态计算，通过结构感知状态融合获取空间上下文，以及使用观测方程进行最终状态计算。我们的理论分析表明，Spatial-Mamba在相同的矩阵乘法框架下统一了原始Mamba和线性注意力，从而提供了对我们方法的更深入理解。实验结果表明，Spatial-Mamba即使在单次扫描的情况下，也能在图像分类、检测和分割方面达到或超过最先进的基于SSM的模型。
🔬 方法详解
问题定义：现有视觉状态空间模型（SSM）在处理图像时，通常将二维图像转换为一维序列，然后通过各种扫描方式来捕捉空间依赖关系。这种方法的痛点在于，一方面，扫描方式难以充分捕捉图像中复杂的空间结构信息；另一方面，较长的扫描路径会显著增加计算复杂度，限制了模型在大规模图像上的应用。
核心思路：Spatial-Mamba的核心思路是直接在状态空间中建立邻域连接，从而避免了将图像转换为一维序列的必要。通过引入结构感知的状态融合机制，模型能够更好地利用图像的空间结构信息，从而提升对图像的理解能力。这种设计旨在克服传统视觉SSM在空间信息建模方面的局限性，并降低计算成本。
技术框架：Spatial-Mamba的处理流程主要包含三个阶段：1) 初始状态计算：通过单向扫描图像，计算每个位置的初始状态；2) 空间上下文获取：利用结构感知的状态融合方程，通过扩张卷积在状态空间中融合邻域信息，从而获取空间上下文；3) 最终状态计算：使用观测方程，基于融合后的状态计算最终的输出。这种三阶段的框架旨在高效地利用图像的空间结构信息，并实现高效的计算。
关键创新：Spatial-Mamba最重要的技术创新点在于其结构感知的状态融合机制。与传统的序列状态转移不同，Spatial-Mamba通过扩张卷积在状态空间中直接融合邻域信息，从而能够更好地捕捉图像的空间结构依赖关系。这种方法避免了将图像转换为一维序列的必要，从而降低了计算复杂度，并提升了模型对空间信息的建模能力。此外，论文还从理论上证明了Spatial-Mamba统一了原始Mamba和线性注意力机制。
关键设计：Spatial-Mamba的关键设计包括：1) 扩张卷积：使用扩张卷积来捕捉不同尺度的空间依赖关系，从而更好地利用图像的空间结构信息；2) 状态融合方程：设计了一种新的状态融合方程，用于在状态空间中融合邻域信息；3) 单向扫描：采用单向扫描来计算初始状态，从而降低计算复杂度。具体的参数设置和网络结构细节可以在论文的实验部分找到。
🖼️ 关键图片
📊 实验亮点
Spatial-Mamba在图像分类、目标检测和语义分割等任务上取得了显著的成果。例如，在ImageNet图像分类任务上，Spatial-Mamba即使使用单次扫描，也能达到或超过现有SOTA的基于SSM的模型。在目标检测和语义分割任务上，Spatial-Mamba也表现出优异的性能，证明了其有效性和泛化能力。
🎯 应用场景
Spatial-Mamba具有广泛的应用前景，可应用于图像分类、目标检测、语义分割等多种计算机视觉任务。其高效的空间信息建模能力使其在处理高分辨率图像和视频时具有优势。未来，Spatial-Mamba有望应用于自动驾驶、医学图像分析、遥感图像处理等领域，提升相关应用的性能和效率。
📄 摘要（原文）
Selective state space models (SSMs), such as Mamba, highly excel at capturing long-range dependencies in 1D sequential data, while their applications to 2D vision tasks still face challenges. Current visual SSMs often convert images into 1D sequences and employ various scanning patterns to incorporate local spatial dependencies. However, these methods are limited in effectively capturing the complex image spatial structures and the increased computational cost caused by the lengthened scanning paths. To address these limitations, we propose Spatial-Mamba, a novel approach that establishes neighborhood connectivity directly in the state space. Instead of relying solely on sequential state transitions, we introduce a structure-aware state fusion equation, which leverages dilated convolutions to capture image spatial structural dependencies, significantly enhancing the flow of visual contextual information. Spatial-Mamba proceeds in three stages: initial state computation in a unidirectional scan, spatial context acquisition through structure-aware state fusion, and final state computation using the observation equation. Our theoretical analysis shows that Spatial-Mamba unifies the original Mamba and linear attention under the same matrix multiplication framework, providing a deeper understanding of our method. Experimental results demonstrate that Spatial-Mamba, even with a single scan, attains or surpasses the state-of-the-art SSM-based models in image classification, detection and segmentation. Source codes and trained models can be found at https://github.com/EdwardChasel/Spatial-Mamba.