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

2024-10-19

arXiv.org

AI Summary

To address the insufficient spatial structure modeling and high computational cost of long-range dependency capture in 2D vision tasks, this paper proposes the Structure-Aware State Space Model (S3M). S3M introduces a novel structure-aware state fusion equation, unifying the theoretical frameworks of Mamba and linear attention. Its three-stage architecture integrates dilated convolutions, unidirectional scanning, and an observation equation to efficiently model pixel-wise neighborhood connectivity in a single scan. Unlike conventional state space models (SSMs), S3M explicitly encodes 2D spatial relationships without increasing scan path length. Extensive experiments demonstrate that S3M consistently outperforms existing SSM-based methods on image classification, object detection, and semantic segmentation—achieving superior efficiency and accuracy in long-range dependency modeling.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Large Multimodal Models (LMMs)
• Intelligent Robots: State Estimation

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search

Abstract

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.

Problem

• Enhances 2D vision task performance
• Reduces computational cost in SSMs
• Captures complex image spatial structures

Innovation

• Structure-aware state fusion
• Dilated convolutions for spatial dependencies
• Unidirectional scan initial state computation