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

## Abstract

Abstract Selective state space models (SSMs), such as Mamba (Gu & Dao, 2023 ) , 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 .

## 1 Introduction

State space models (SSMs) are powerful tools for analyzing dynamic systems with hidden states, and they have long been utilized in fields like control theory, signal processing, and economics (Friston et al., ; Hafner et al., ; Gu et al., ). SSMs have been recently introduced into deep learning (Gu et al., ), especially in natural language processing (NLP) (Gu & Dao, ), thanks to their use of specially parameterized matrices. A major advancement is the introduction of selective mechanisms and hardware-aware optimizations for parallel computing, as demonstrated by Mamba (Gu & Dao, ), which selectively retains or discards information based on the relevance of each element in a sequence, efficiently modeling long-distance dependencies with linear complexity.

The significant success of Mamba in NLP inspires researchers to investigate how SSMs can be applied to visual tasks. Unlike 1D sequences, visual data are typically characterized by 2D spatial structures.
Therefore, it is crucial to maintain the spatial dependencies within images while adapting the information selection and propagation mechanisms in Mamba.
Existing visual SSMs (Zhu et al., ; Liu et al., ; Yang et al., ; Huang et al., ; He et al., ; Xiao et al., ) often use some scanning strategies to flatten 2D visual data into several 1D sequences from different directions, and then process the flattened 1D sequences using the original Mamba.
These scanning strategies can be broadly categorized into three types: sweeping scan, continuous scan and local scan, as shown in Figs. 1(a)-1(c), respectively.

Vim (Zhu et al., ) and VMamba (Liu et al., ) employ sweeping bidirectional scan and four-way scan, which are illustrated in Fig. 1(a). These strategies aim to reduce spatial direction sensitivity and adapt the network architecture for visual tasks.
Yang et al. () argued that sweeping scans neglect the importance of spatial continuity, and they introduced a continuous scanning order, as shown in Fig. 1(b), to better integrate the inductive biases from visual data.
Huang et al. () argued that scanning the entire image may not effectively capture local spatial relationships. Instead, they presented LocalMamba by using several local scanning modes, as shown in Fig. 1(c), which divides an image into distinct windows to capture local dependencies.
Beyond the above mentioned scanning patterns, other scanning methods such as Hilbert scanning (He et al., ) and dynamic tree scanning (Xiao et al., ) have also been proposed to adapt SSMs to visual tasks.

Figure: (a) Sweeping Scan
Refer to caption: x1.png

While the existing scanning strategies partially address the issue of aligning spatial structures with sequential SSMs, they are limited in model effectiveness and efficiency.
On one hand, directional scanning inevitably alters the spatial relationships between pixels, disrupting the inherent spatial context in an image. For example, in the sweeping scan (Fig. 1(a)), the distance between a pixel and its left or right neighbor is 1, while its distance to the top or bottom neighbor equals to the image width. This distortion can hinder the models to understand the spatial relationships in the original visual data.
On the other hand, the fixed scanning paths, such as the commonly used four-directional scan (Figs. 1(a) and  1(b)), are not effective enough to capture the complex and varying spatial relationships in an image, while introducing more scanning directions would also result in excessive computations.
Therefore, it is imperative to explore how to design more effective and structure-aware SSMs for visual tasks.

To achieve this goal, we make an important attempt in this paper and present Spatial-Mamba, which is designed to capture the spatial dependencies of neighboring features in the latent state space. The processing flow of Spatial-Mamba consists of three stages.
First, as shown in the left of Fig. 1(d), visual data are converted into sequential data using unidirectional sweeping scan. The state variables are computed based on the state transition equations of the original SSMs and then reshaped back into the visual format.
Second, the state variables are processed through a structure-aware state fusion (SASF) equation, which employs dilated convolutions to re-weight and merge nearby state variables, as shown in the left of Fig. 1(d).
Finally, these structure-aware state variables are fed into the observation equation to produce the final output variables. The SASF equation not only enables efficient skip connections between non-sequential elements in sequences but also enhances the model ability to capture spatial relationships, leading to more accurate representations of the underlying visual structure.
Furthermore, we show that Spatial-Mamba, original Mamba and linear attention can all be represented under the same framework using structured matrices, which offers a more coherent understanding of our proposed method.
We validate the superiority of Spatial-Mamba across fundamental vision tasks such as image classification, detection, and segmentation. The results demonstrate that Spatial-Mamba, even with a single scan, achieves or surpasses the performance of recent state-of-the-arts using different scanning strategies.

## 2 Related Work

State space models (SSMs).
Gu et al. () firstly introduced the linear state space layer (LSSL) into the HiPPO framework (Gu et al., ) to efficiently handle the long-range dependencies in long sequences.
Gu et al. () then significantly improved the efficiency of SSMs by representing the parameters as diagonal plus low-rank matrix. The so-called S4 model triggers a wave of structured SSMs (Smith et al., ; Fu et al., ; Gupta et al., ; Gu & Dao, ).
Smith et al. () proposed S5 by introducing parallel scans to S4 layer while maintaining the computational efficiency of S4.
Recently, Gu & Dao () developed Mamba, which incorporates a data-dependent selection mechanism into S4 layer and simplifies the computation and architecture in a hardware-friendly way, achieving Transformer-like modeling capability with linear complexity.
Building on that, Dao & Gu () presented Mamba2, which reveals the connections between SSMs and attention with specific structured matrix. This framework simplifies the parameter matrix to a scalar representation, making it feasible to explore larger and more expressive state spaces without sacrificing efficiency.

Visual SSMs.
Although traditional SSMs perform well in processing NLP sequential data and capturing temporal dependencies, they struggle in handling multi-dimensional spatial structure inherent in visual data. This limitation poses a challenge for developing effective visual SSMs. S4ND (Nguyen et al., ) is among the first SSM-based models for multi-dimensional data, which separates each dimension with an independent 1D SSM. Baron et al. () generalized S4ND as a discrete multi-axial system and proposed the 2D-SSM spatial layer, successfully extending the 1D SSMs to 2D SSMs. More recent visual SSMs prefer to design multiple scanning orders or patterns to maintain the spatial consistency, including bidirectional (Liu et al., ), four-way (Zhu et al., ), continuous (Yang et al., ), zigzag (Hu et al., ), window-based (Huang et al., ), and topology-based scanning (He et al., ; Xiao et al., ).
These visual SSMs have been used in multimodal foundation models (Qiao et al., ; Mo & Tian, ), image restoration (Guo et al., ; Shi et al., ), medical image analysis (Yue & Li, ; Ma et al., ; He et al., ; Liao et al., ) and other visual tasks (Chen et al., ; Li et al., ; Yao et al., ), demonstrating the potential of SSMs in visual data understanding.

## 3 Preliminaries

SSMs are commonly used for analyzing sequential data and modeling continuous linear time-invariant (LTI) systems (Williams & Lawrence, ). An input sequence $\displaystyle u(t)\in\mathbb{R}$ is transformed into an output sequence $\displaystyle y(t)\in\mathbb{R}$ through a state variable $\displaystyle x(t)\in\mathbb{C}^{N}$. Here, $t>0$ represents the time index, and $N$ indicates the dimension of the state variable. This dynamic system can be described by the linear state transition and observation equations (Kalman, ):
$\displaystyle x^{\prime}(t)={\bm{A}}x(t)+{\bm{B}}u(t),y(t)={\bm{C}}x(t)+{\bm{D
}}u(t)$,
where $\displaystyle{\bm{A}}\in\mathbb{C}^{N\times N}$ is the state transition matrix, $\displaystyle{\bm{B}}$, $\displaystyle{\bm{C}}\in\mathbb{C}^{N}$ and $\displaystyle{\bm{D}}\in\mathbb{C}^{1}$ control the dynamics of the system. The state transition and observation equations describe how the system evolves over time and how the state variables relate to the observed outputs.

To effectively integrate continuous-time SSMs into the deep learning framework, it is essential to discretize the continuous-time models. One commonly employed technique is the Zero-Order Hold (ZOH) discretization (Gu & Dao, ). The ZOH method approximates the continuous-time system by holding the input constant over each discrete time interval. Specifically, given a timescale $\mathbf{\Delta}$, which represents the interval between discrete time steps, and defining $\overline{{\bm{A}}}$ and $\overline{{\bm{B}}}$ as discrete parameters, the ZOH rule is applied as $\displaystyle\overline{{\bm{A}}}=e^{\mathbf{\Delta}{\bm{A}}}$ and $\overline{{\bm{B}}}=(\mathbf{\Delta}{\bm{A}})^{-1}(e^{\mathbf{\Delta}{\bm{A}}}
-I)\mathbf{\Delta}{\bm{B}}$.

However, real-world processes often change over time and cannot be accurately described by a LTI system. As highlighted in Mamba (Gu & Dao, ), time-varying systems can focus more on relevant information and offer a more accurate and realistic representation of dynamic systems.
In Mamba, the parameters of the SSMs are made context-aware and adaptive through selective functions. This is achieved by modifying the parameters $\mathbf{\Delta},{\bm{B}},{\bm{C}}$ as simple functions of the input sequence $u_{t}$, resulting in input-dependent parameters $\mathbf{\Delta}_{t}=s_{\Delta}(u_{t}),{\bm{B}}_{t}=s_{B}(u_{t})$ and ${\bm{C}}_{t}=s_{C}(u_{t})$. Then the input-dependent discrete parameters $\overline{{\bm{A}}}_{t}$ and $\overline{{\bm{B}}}_{t}$ can be calculated accordingly. Consequently, the discrete state transition and observation equations can be calculated as follows:

$$ $\displaystyle x_{t}=\overline{{\bm{A}}}_{t}x_{t-1}+\overline{{\bm{B}}}_{t}u_{t },\quad y_{t}={\bm{C}}_{t}x_{t}+{\bm{D}}u_{t}.$ (1) $$

A simplified illustration of the above process is shown in Fig. 2(a).

Figure: (a) SSM in Mamba
Refer to caption: x5.png

Figure: (a) Input
Refer to caption: x7.png

## 4 Spatial-Mamba

### 4.1 Formulation of Spatial-Mamba

Spatial-Mamba is designed to capture the spatial dependencies of neighboring features in the latent state space. To achieve this goal, different from previous methods (Zhu et al., ; Liu et al., ; Huang et al., ) that commonly employ multiple scanning directions, we introduce a new structure-aware state fusion (SASF) equation into the original Mamba formulas (refer to Eq. (1)). The entire process of Spatial-Mamba can be described by three equations: the state transition equation, the SASF and the observation equation, which are formulated as:

$$ $\displaystyle x_{t}=\overline{{\bm{A}}}_{t}x_{t-1}+\overline{{\bm{B}}}_{t}u_{t },\quad h_{t}=\sum_{k\in\Omega}\alpha_{k}x_{\rho_{k}(t)},\quad y_{t}={\bm{C}}_ {t}h_{t}+{\bm{D}}u_{t},$ (2) $$

where $x_{t}$ is the original state variable, $h_{t}$ is the structure-aware state variable, $\Omega$ is the neighbor set, and $\rho_{k}(t)$ is the index of the $k$-th neighbor of position $t$. Fig. 2(b) illustrates the SSM flow in the proposed Spatial-Mamba. Compared with the original Mamba in Fig. 2(a), we can see that the original state variable $x_{t}$ is directly influenced by its previous state $x_{t-1}$, while the structure-aware state variable $h_{t}$ incorporates additional neighboring state variables $x_{\rho_{1}(t)},x_{\rho_{2}(t)},\dots,x_{\rho_{K}(t)}$ through a fusion mechanism, where $K=|\Omega|$ denotes the size of neighbor set. By considering both temporal and spatial information, the fused state variable $h_{t}$ gains a richer context, leading to improved adaptability and a more comprehensive understanding of the image.

The proposed Spatial-Mamba can be implemented in three steps. As shown in Fig. 1(d), the input image is first flattened into a 1D sequence $u_{t}$, with which the state $x_{t}$ is computed using the state transition equation $x_{t}=\overline{{\bm{A}}}_{t}x_{t-1}+\overline{{\bm{B}}}_{t}u_{t}$. The computed states are then reshaped back into the 2D format.
To enable each state to be aware of its neighboring states in the 2D space, we introduce the SASF equation $h_{t}=\sum_{k\in\Omega}\alpha_{k}x_{\rho_{k}(t)}$. For a state variable $x_{t}$, we apply linear weighting to its neighboring states $\rho_{k}(t)$ in the neighborhood $\Omega$ using weights $\alpha_{k}$ so that we can effectively integrate local dependency information into a new state $h_{t}$, resulting in a more contextually informative representation. Finally, the output is generated from this enriched state $h_{t}$ through the observation equation $y_{t}={\bm{C}}_{t}h_{t}+{\bm{D}}u_{t}$. This SASF approach helps the model to incorporate the local structural information in visual learning while retaining the benefits of original Mamba.

In practice, we simply employ multi-scale dilated convolutions (Yu, ) to linearly weight adjacent state variables, enhancing spatial relationship characterizations and enabling skip connections. Specifically, we use three $3\times 3$ depth-wise filters with dilation factors $d=1,3,5$ to construct the neighbor set $\Omega_{d}=\{(i,j)|i,j\in\{-d,0,d\}\}$. The SASF equation in Eq. (2) can be rewritten as:

$$ $h_{t}=\sum_{d=1,3,5}\sum_{i,j\in\Omega_{d}}k_{ij}^{d}\cdot x_{t+iw+j},$ (3) $$

where $k_{ij}^{d}$ represents the filter weight for dilation factor $d$ at position $(i,j)$, $x_{t+iw+j}$ denotes the neighbor of the state $x_{t}$ located at the position $(i,j)$, and $w$ indicates the width of the image.

To gain an intuitive understanding of SASF, we visualize the original state variables and our proposed structure-aware state variables in Fig. 3. One can see that the original state $x_{t}$ (Fig. 3(b)) struggles to differentiate between the foreground and background. In contrast, the structure-aware state $h_{t}$ (Fig. 3(c)), which has been refined through a SASF process, effectively separates these regions. Moreover, the original state variable $x_{t}$ in Fig. 3(d) only shows horizontal attenuation along the scanning direction (gradually darkening from the brightest value in the upper left corner), while the fused state variable $h_{t}$ in Fig. 3(e) demonstrates decay along the horizontal, vertical and diagonal directions. This improvement stems from its ability to leverage spatial relationships within the image, leading to a more accurate and context-aware representation.

### 4.2 Network Architecture

The overall architecture of Spatial-Mamba is depicted in Fig. 4. It consists of four successive stages, resembling the structure of Swin-Transformer (Liu et al., ). We introduce three variants of Spatial-Mamba model at different scales: Spatial-Mamba-T (tiny), Spatial-Mamba-S (small), and Spatial-Mamba-B (base). The detailed configurations are provided in Appendix A. Specifically, an input image $I\in\mathbb{R}^{H\times W\times 3}$ is first processed by an overlapped stem layer to generate a 2D feature map with dimension of $\frac{H}{4}\times\frac{W}{4}\times C$. This feature map is then fed into four successive stages. Each stage comprises multiple Spatial-Mamba blocks, followed by a down-sampling layer with a factor of 2 (except for the last stage), resulting in hierarchical features. Finally, a head layer is employed to process these features to produce corresponding image representations for specific downstream tasks.

The Spatial-Mamba block forms the fundamental building unit of our architecture, which consists of a Structure-aware SSM and a Feed-Forward Network (FFN) with skip connections, as illustrated in the bottom left of Fig. 4. Building upon the Mamba block design (Gu & Dao, ), the Structure-aware SSM, illustrated in the bottom right of Fig. 4, is implemented by substituting the original 1D causal convolution with a $3\times 3$ depth-wise convolution and replacing the original S6 module with our proposed SASF module, achieving local neighborhood connectivity in state spaces with linear complexity. Moreover, a local perception unit (LPU) (Guo et al., ) is employed before the Spatial-Mamba block and FFN to extract local information inside image patches.

Figure: Figure 4: Overall network architecture of Spatial-Mamba.
Refer to caption: x13.png

### 4.3 Connection with Original Mamba and Linear Attention

We analyze in-depth the similarities and disparities among linear attention (Katharopoulos et al., ), original Mamba (Gu & Dao, ) and our Spatial-Mamba, providing a better understanding of the working mechanism of our proposed method. Detailed derivations are provided in Appendix B.

Linear attention is an improved self-attention (SA) mechanism, reducing the computational complexity of SA to linear by using a kernel function $\phi$. For an input sequence $u_{t}$, the query $q_{t}$, key $k_{t}$, and value $v_{t}$ are computed by projecting $u_{t}$ with different weight matrices. In autoregressive models, to prevent the model from attending to future tokens, the $t$-th query is restricted by the previous keys, *i.e*., $k_{s},s\leq t$. Thus, if the kernel function $\phi$ is an identity map, the calculation of single head linear attention without normalization can be formulated as: $y_{t}=\sum_{s\leq t}q_{t}(k_{s}^{T}v_{s})$. Letting $x_{t}=\sum_{s\leq t}k_{s}^{T}v_{s}$, then we have $x_{t}=x_{t-1}+k_{t}^{T}v_{t}$. The linear attention can be rewritten as $y_{t}=q_{t}x_{t}$. This reveals that linear attention is actually a special case of linear recursion. The hidden state variable $x_{t}$ is updated by the outer-product $k_{s}^{T}v_{s}$, and the final output $y_{t}$ is observed by multiplying $x_{t}$ with the query $q_{t}$. If we define ${\bm{C}}_{t}=q_{t}$ and $\overline{{\bm{B}}}_{s}=k_{s}^{T}$, the linear attention can be expressed in a form similar to that of SSM:

$$ $y_{t}=\sum_{s\leq t}{\bm{C}}_{t}\overline{{\bm{B}}}_{s}v_{s},$ (4) $$

where $v_{s}$ is a linear transformation of $u_{s}$, *i.e*., $v_{s}=\text{Linear}(u_{s}).$

Mamba is essentially defined in Eq. (1).
By setting the initial state variable $x_{0}$ to zero, the state variables can be derived recursively as $x_{t}=\sum_{s\leq t}\overline{{\bm{A}}}_{s:t}^{\times}\overline{{\bm{B}}}_{s}u
_{s}$. Here, $\overline{{\bm{A}}}_{s:t}^{\times}:=\Pi_{i=s+1}^{t}\overline{{\bm{A}}}_{i}$
denotes the product of the state transition matrices with indices from $s+1$ to $t$ for $s<t$, and its value is $1$ when $s=t$. The final output of the observation equation without ${\bm{D}}u_{t}$ can be rewritten as:

$$ $y_{t}=\sum_{s\leq t}{\bm{C}}_{t}\ \overline{{\bm{A}}}_{s:t}^{\times}\ \overline{{\bm{B}}}_{s}u_{s}.$ (5) $$

Spatial-Mamba. Based on Eq. (2) and Eq. (5), the structure-aware state variables from our proposed SASF equation can be expressed as
$h_{t}=\sum_{k\in\Omega}\ \sum_{s\leq\rho_{k}(t)}\alpha_{k}\overline{{\bm{A}}}_
{s:{\rho_{k}(t)}}^{\times}\ \overline{{\bm{B}}}_{s}u_{s}$. By omitting $Du_{t}$ for simplicity of expression, the final output of Spatial-Mamba can be reformulated as follows:

$$ $y_{t}=\sum_{k\in\Omega}\sum_{s\leq\rho_{k}(t)}\alpha_{k}\ {\bm{C}}_{t}\ \overline{{\bm{A}}}_{s:\rho_{k}(t)}^{\times}\ \overline{{\bm{B}}}_{s}u_{s}.$ (6) $$

Figure: (a) Linear attention
Refer to caption: x14.png

Remarks. From the above analysis, we can see that all the three paradigms — linear attention, Mamba, and Spatial-Mamba — can be modeled within a unified matrix multiplication framework, specifically $y={\bm{M}}u$. The differences lie in the structure of ${\bm{M}}$. For both linear attention and Mamba, ${\bm{M}}$ takes the form of a lower triangular matrix, whereas for Spatial-Mamba, ${\bm{M}}$ is an adjacency matrix.
Fig. 5 provides a visualization of these matrices.
In linear attention, the positions of brighter values remain consistent along the vertical direction, which indicates that the SA mechanism puts its focus on a small set of image tokens. Mamba, on the other hand, shows a decaying pattern over time, which is attributed to the influence of its state transition matrix $\overline{{\bm{A}}}_{t}$. This dynamic transition allows Mamba to shift its focus among previous image tokens. Unlike linear attention and Mamba, our Spatial-Mamba considers the weighted summation of all states within a broader spatial neighborhood $\Omega$, allowing for a more comprehensive representation of spatial relationships.

## 5 Experimental Results

In this section, we conduct a series of experiments to compare Spatial-Mamba with leading benchmark models, including those based on Convolutional Neural Networks (CNNs) (Radosavovic et al., ; Liu et al., ; Yu & Wang, ), Vision Transformers (Dosovitskiy et al., ; Liu et al., ; Wang et al., ; Hassani et al., ), and recent visual SSMs (Nguyen et al., ; Zhu et al., ; Liu et al., ; Huang et al., ). Following previous works (Liu et al., ; ), we train three variants of Spatial-Mamba, namely Spatial-Mamba-T, Spatial-Mamba-S and Spatial-Mamba-B. The detailed configurations are provided in Appendix A. The performance evaluation is conducted on fundamental visual tasks, including image classification, object detection, and semantic segmentation.

### 5.1 Image Classification on ImageNet-1K

**Table 1: Comparison of classification performance on ImageNet-1K, where ‘Throughput’ is measured using an A100 GPU with an input resolution of $224\times 224$.**
| Arch. | Method | Im. size | #Param. (M) | FLOPs (G) | Throughput$\uparrow$ | Top-1 acc.$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- |
| CNN | RegNetY-4G | $224^{2}$ | 21M | 4.0G | - | 80.0 |
| RegNetY-8G | $224^{2}$ | 39M | 8.0G | - | 81.7 |  |
| RegNetY-16G | $224^{2}$ | 84M | 16.0G | - | 82.9 |  |
| ConvNeXt-T | $224^{2}$ | 29M | 4.5G | 1189 | 82.1 |  |
| ConvNeXt-S | $224^{2}$ | 50M | 8.7G | 682 | 83.1 |  |
| ConvNeXt-B | $224^{2}$ | 89M | 15.4G | 435 | 83.8 |  |
| Transformer | ViT-B/16 | $384^{2}$ | 86M | 55.4G | - | 77.9 |
| ViT-L/16 | $384^{2}$ | 307M | 190.7G | - | 76.5 |  |
| DeiT-S | $224^{2}$ | 22M | 4.6G | 1759 | 79.8 |  |
| DeiT-B | $224^{2}$ | 86M | 17.5G | 500 | 81.8 |  |
| DeiT-B | $384^{2}$ | 86M | 55.4G | 498 | 83.1 |  |
| Swin-T | $224^{2}$ | 28M | 4.5G | 1244 | 81.3 |  |
| Swin-S | $224^{2}$ | 50M | 8.7G | 718 | 83.0 |  |
| Swin-B | $224^{2}$ | 88M | 15.4G | 458 | 83.5 |  |
|  | NAT-T | $224^{2}$ | 28M | 4.3G | - | 83.2 |
|  | NAT-S | $224^{2}$ | 51M | 7.8G | - | 83.0 |
|  | NAT-B | $224^{2}$ | 90M | 13.7G | - | 84.3 |
| SSM | S4ND-ConvNeXt-T | $224^{2}$ | 30M | - | 683 | 82.2 |
| S4ND-ViT-B | $224^{2}$ | 89M | - | 397 | 80.4 |  |
| ViM-S | $224^{2}$ | 26M | - | 811 | 80.5 |  |
| VMamba-T | $224^{2}$ | 30M | 4.9G | 1686 | 82.6 |  |
| VMamba-S | $224^{2}$ | 50M | 8.7G | 877 | 83.6 |  |
| VMamba-B | $224^{2}$ | 89M | 15.4G | 646 | 83.9 |  |
| LocalVMamba-T | $224^{2}$ | 26M | 5.7G | 394 | 82.7 |  |
| LocalVMamba-S | $224^{2}$ | 50M | 11.4G | 227 | 83.7 |  |
| Spatial-Mamba-T | $224^{2}$ | 27M | 4.5G | 1438 | 83.5 |  |
| Spatial-Mamba-S | $224^{2}$ | 43M | 7.1G | 988 | 84.6 |  |
| Spatial-Mamba-B | $224^{2}$ | 96M | 15.8G | 665 | 85.3 |  |

Settings. We first evaluate the representation learning capabilities of Spatial-Mamba in image classification on ImageNet-1K (Deng et al., ). We adopted the experimental configurations used in previous works (Liu et al., ; ), which are detailed in Appendix A.
We compare our method with state-of-the-art approaches, including RegNetY (Radosavovic et al., ), ConvNeXt (Liu et al., ), ViT (Dosovitskiy et al., ), DeiT (Touvron et al., ), Swin (Liu et al., ), NAT (Hassani et al., ), S4ND (Nguyen et al., ), Vim (Zhu et al., ), VMamba (Liu et al., ), and LocalVMamba (Huang et al., ).

**Table 2: Comparison of object detection and instance segmentation performance on COCO with Mask R-CNN (He et al., ) detector. FLOPs are calculated with input resolution of $1280\times 800$.**
| Mask R-CNN 1$\times$ schedule |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Backbone | AP${}^{\text{b}}\uparrow$ | AP${}^{\text{b}}_{50}$$\uparrow$ | AP${}^{\text{b}}_{75}$$\uparrow$ | AP${}^{\text{m}}\uparrow$ | AP${}^{\text{m}}_{50}$$\uparrow$ | AP${}^{\text{m}}_{75}$$\uparrow$ | #Param. | FLOPs |
| ResNet-50 | 38.2 | 58.8 | 41.4 | 34.7 | 55.7 | 37.2 | 44M | 260G |
| Swin-T | 42.7 | 65.2 | 46.8 | 39.3 | 62.2 | 42.2 | 48M | 267G |
| ConvNeXt-T | 44.2 | 66.6 | 48.3 | 40.1 | 63.3 | 42.8 | 48M | 262G |
| PVTv2-B2 | 45.3 | 66.1 | 49.6 | 41.2 | 64.2 | 44.4 | 45M | 309G |
| ViT-Adapter-S | 44.7 | 65.8 | 48.3 | 39.9 | 62.5 | 42.8 | 48M | 403G |
| MambaOut-T | 45.1 | 67.3 | 49.6 | 41.0 | 64.1 | 44.1 | 43M | 262G |
| VMamba-T | 47.3 | 69.3 | 52.0 | 42.7 | 66.4 | 45.9 | 50M | 271G |
| LocalVMamba-T | 46.7 | 68.7 | 50.8 | 42.2 | 65.7 | 45.5 | 45M | 291G |
| Spatial-Mamba-T | 47.6 | 69.6 | 52.3 | 42.9 | 66.5 | 46.2 | 46M | 261G |
| ResNet-101 | 38.2 | 58.8 | 41.4 | 34.7 | 55.7 | 37.2 | 63M | 336G |
| Swin-S | 44.8 | 68.6 | 49.4 | 40.9 | 65.3 | 44.2 | 69M | 354G |
| ConvNeXt-S | 45.4 | 67.9 | 50.0 | 41.8 | 65.2 | 45.1 | 70M | 348G |
| PVTv2-B3 | 47.0 | 68.1 | 51.7 | 42.5 | 65.2 | 45.7 | 63M | 397G |
| MambaOut-S | 47.4 | 69.1 | 52.4 | 42.7 | 66.1 | 46.2 | 65M | 354G |
| VMamba-S | 48.7 | 70.0 | 53.4 | 43.7 | 67.3 | 47.0 | 70M | 349G |
| LocalVMamba-S | 48.4 | 69.9 | 52.7 | 43.2 | 66.7 | 46.5 | 69M | 414G |
| Spatial-Mamba-S | 49.2 | 70.8 | 54.2 | 44.0 | 67.9 | 47.5 | 63M | 315G |
| Swin-B | 46.9 | - | - | 42.3 | 66.3 | 46.0 | 88M | 496G |
| ConvNeXt-B | 47.0 | 69.4 | 51.7 | 42.7 | 66.3 | 46.0 | 107M | 486G |
| PVTv2-B5 | 47.4 | 68.6 | 51.9 | 42.5 | 65.7 | 46.0 | 102M | 557G |
| ViT-Adapter-B | 47.0 | 68.2 | 51.4 | 41.8 | 65.1 | 44.9 | 102M | 557G |
| MambaOut-B | 47.4 | 69.3 | 52.2 | 43.0 | 66.4 | 46.3 | 100M | 495G |
| VMamba-B | 49.2 | 71.4 | 54.0 | 44.1 | 68.3 | 47.7 | 108M | 485G |
| Spatial-Mamba-B | 50.4 | 71.8 | 55.3 | 45.1 | 69.1 | 49.1 | 115M | 494G |
| Mask R-CNN 3$\times$ MS schedule |  |  |  |  |  |  |  |  |
| Backbone | AP${}^{\text{b}}\uparrow$ | AP${}^{\text{b}}_{50}$$\uparrow$ | AP${}^{\text{b}}_{75}$$\uparrow$ | AP${}^{\text{m}}\uparrow$ | AP${}^{\text{m}}_{50}$$\uparrow$ | AP${}^{\text{m}}_{75}$$\uparrow$ | #Param. | FLOPs |
| Swin-T | 46.0 | 68.1 | 50.3 | 41.6 | 65.1 | 44.9 | 48M | 267G |
| ConvNeXt-T | 46.2 | 67.9 | 50.8 | 41.7 | 65.0 | 44.9 | 48M | 262G |
| NAT-T | 47.7 | 69.0 | 52.6 | 42.6 | 66.1 | 45.9 | 48M | 258G |
| VMamba-T | 48.8 | 70.4 | 53.5 | 43.7 | 67.4 | 47.0 | 50M | 271G |
| LocalVMamba-T | 48.7 | 70.1 | 53.0 | 43.4 | 67.0 | 46.4 | 45M | 291G |
| Spatial-Mamba-T | 49.3 | 70.7 | 54.3 | 43.6 | 67.6 | 46.9 | 46M | 261G |
| Swin-S | 48.2 | 69.8 | 52.8 | 43.2 | 67.0 | 46.1 | 69M | 354G |
| ConvNeXt-S | 47.9 | 70.0 | 52.7 | 42.9 | 66.9 | 46.2 | 70M | 348G |
| NAT-S | 48.4 | 69.8 | 53.2 | 43.2 | 66.9 | 46.5 | 70M | 330G |
| VMamba-S | 49.9 | 70.9 | 54.7 | 44.2 | 68.2 | 47.7 | 70M | 349G |
| LocalVMamba-S | 49.9 | 70.5 | 54.4 | 44.1 | 67.8 | 47.4 | 69M | 414G |
| Spatial-Mamba-S | 50.5 | 71.5 | 55.5 | 44.6 | 68.7 | 47.8 | 63M | 315G |

Results. Tab. 1 presents a comprehensive comparison between Spatial-Mamba against state-of-the-art methods. Notably, Spatial-Mamba-T achieves a top-1 accuracy of 83.5%, outperforming the CNN-based ConvNext-T by 1.4% with similar amount of parameters and GFLOPs. Compared to Transformer-based methods, Spatial-Mamba-T exceeds Swin-T by 2.2% and NAT-T by 0.3%. In comparison with SSM-based methods, Spatial-Mamba-T outperforms VMamba-T by 1.0% and LocalVMamba-T by 0.8%.
For other variants, Spatial-Mamba also shows advantages. Specifically, Spatial-Mamba-S and Spatial-Mamba-B achieve top-1 accuracies of 84.6% and 85.3%, respectively, surpassing NAT-S and NAT-B by margins of 1.6% and 1.0%, and VMamba-S and VMamba-B by 1.0% and 1.4%. While Spatial-Mamba-T is slightly slower than VMamba-T due to architectural differences, the Small and Base variants of Spatial-Mamba are faster than their VMamba counterparts. Moreover, both of them are significantly faster than CNN and Transformer-based methods.

### 5.2 Object Detection and Instance Segmentation on COCO

Settings. We evaluate Spatial-Mamba in object detection and instance segmentation tasks using COCO 2017 dataset (Lin et al., ) and MMDetection library (Chen et al., ). We adopt Mask R-CNN (He et al., ) as detector head, apply Spatial-Mamba-T/S/B pre-trained on ImageNet-1K as backbones. Following common practices (Liu et al., ; ), we fine-tune the pre-trained models for 12 epochs (1$\times$ schedule) and 36 epochs with multi-scale inputs (3$\times$ schedule). During training, AdamW optimizer is adopted with an initial learning rate of 0.0001 and a batch size of 16.

Results. Detailed results on COCO are reported in Tab. 2. It can be seen that all variants of Spatial-Mamba outperform their competitors under different schedules. For 1$\times$ schedule, Spatial-Mamba-T achieves a box mAP of 47.6 and a mask mAP of 42.9, surpassing Swin-T/VMamba-T by 4.9/0.3 in box mAP and 3.6/0.2 in mask mAP with fewer parameters and FLOPs, respectively. Similarly, Spatial-Mamba-S/B demonstrate superior performance to other methods under the same configuration. Furthermore, these trends of improved performance hold with the 3$\times$ multi-scale training schedule. Notably, Spatial-Mamba-S achieves the highest box mAP of 50.5 and mask mAP of 44.6, surpassing VMamba-S with a considerable gain of 0.6 and 0.4, respectively.

### 5.3 Semantic Segmentation on ADE20K

**Table 3: Comparison of semantic segmentation on ADE20K with UPerNet (Xiao et al., ) segmentor. FLOPs are calculated with input resolution of $512\times 2048$. ‘SS’ and ‘MS’ represent single-scale and multi-scale testing, respectively.**
| Method | Crop size | mIoU (SS) $\uparrow$ | mIoU (MS) $\uparrow$ | #Param. | FLOPs |
| --- | --- | --- | --- | --- | --- |
| DeiT-S + MLN | $512^{2}$ | 43.1 | 43.8 | 58M | 1217G |
| Swin-T | $512^{2}$ | 44.4 | 45.8 | 60M | 945G |
| ConvNeXt-T | $512^{2}$ | 46.0 | 46.7 | 60M | 939G |
| NAT-T | $512^{2}$ | 47.1 | 48.4 | 58M | 934G |
| MambaOut-T | $512^{2}$ | 47.4 | 48.6 | 54M | 938G |
| VMamba-T | $512^{2}$ | 48.0 | 48.8 | 62M | 949G |
| LocalVMamba-T | $512^{2}$ | 47.9 | 49.1 | 57M | 970G |
| Spatial-Mamba-T | $512^{2}$ | 48.6 | 49.4 | 57M | 936G |
| DeiT-B + MLN | $512^{2}$ | 45.5 | 47.2 | 144M | 2007G |
| Swin-S | $512^{2}$ | 47.6 | 49.5 | 81M | 1039G |
| ConvNeXt-S | $512^{2}$ | 48.7 | 49.6 | 82M | 1027G |
| NAT-S | $512^{2}$ | 48.0 | 49.5 | 82M | 1010G |
| MambaOut-S | $512^{2}$ | 49.5 | 50.6 | 76M | 1032G |
| VMamba-S | $512^{2}$ | 50.6 | 51.2 | 82M | 1028G |
| LocalVMamba-S | $512^{2}$ | 50.0 | 51.0 | 81M | 1095G |
| Spatial-Mamba-S | $512^{2}$ | 50.6 | 51.4 | 73M | 992G |
| Swin-B | $512^{2}$ | 48.1 | 49.7 | 121M | 1188G |
| ConvNeXt-B | $512^{2}$ | 49.1 | 49.9 | 122M | 1170G |
| NAT-B | $512^{2}$ | 48.5 | 49.7 | 123M | 1137G |
| MambaOut-B | $512^{2}$ | 49.6 | 51.0 | 112M | 1178G |
| VMamba-B | $512^{2}$ | 51.0 | 51.6 | 122M | 1170G |
| Spatial-Mamba-B | $512^{2}$ | 51.8 | 52.6 | 127M | 1176G |

Settings. To assess the performance of Spatial-Mamba on semantic segmentation task, we train our models with the widely used UPerNet segmentor (Xiao et al., ) and MMSegmenation toolkit (Contributors, ). Consistent with previous work (Liu et al., ; ), we pre-train our model on ImageNet-1K, and use it as the backbone to train UPerNet on ADE20K dataset (Zhou et al., ). This training process encompasses 160K iterations with a batch size of 16. The AdamW is used as the optimizer with a weight decay of 0.01. The learning rate is set to $6\times 10^{-5}$ with a linear learning rate decay. All the input images are cropped into $512\times 512$.

Results. The results on semantic segmentation are summarized in Tab. 3. Spatial-Mamba variants consistently achieve impressive performance. For instance, Spatial-Mamba-T attains a single-scale mIoU of 48.6 and a multi-scale mIoU of 49.4. This signifies an improvement of 1.5 mIoU over NAT-T and 0.6 mIoU over VMamba-T in single-scale input. This advantage is maintained with multi-scale input, where Spatial-Mamba-T is 1.0 mIoU higher than NAT-T and 0.6 mIoU higher than VMamba-T. Furthermore, Spatial-Mamba-B achieves the best performance with a multi-scale mIoU of 52.6.

### 5.4 Ablation Studies

In this section, we ablate various key components of Spatial-Mamba-T on ImageNet-1K classification task in Tab. 4. Based on the configurations of Swin-T (Liu et al., ), we construct the baseline model as Spatial-Mamba-T but without the SASF module. This baseline uses a $4\times 4$ convolution with a stride of $4$ as stem layer and merges $2\times 2$ neighboring patches for down-sampling.

Neighbor set. First, adjustments to the neighbor set $\Omega$ reveal that increasing the size from a $3\times 3$ neighbourhood to $5\times 5$ results in a gradual improvement in accuracy, from 82.3% to 82.5%, albeit with a corresponding decrease in throughput. Furthermore, employing a broader dilated neighbor set with factors $d=1,3,5$ increases the accuracy to 82.7%, while reducing throughput to 1158. This suggests a trade-off between speed and larger neighbor set.

**Table 4: Ablation studies on Spatial-Mamba-T for neighbor set $\Omega$ and other module designs.**
| Module design | #Param. (M) | FLOPs (G) | Throughput $\uparrow$ | Top-1 acc.$\uparrow$ |
| --- | --- | --- | --- | --- |
| Baseline | 25M | 4.4G | 1706 | 82.0 |
| $\Omega=\{3\times 3\}$ | 25M | 4.4G | 1557 | 82.3 |
| $\Omega=\{5\times 5\}$ | 25M | 4.5G | 1461 | 82.5 |
| $\Omega=\left\{\Omega_{d}|d=1,3,5\right\}$ | 25M | 4.5G | 1209 | 82.7 |
| $+$ Overlapped Stem | 27M | 4.5G | 1158 | 82.9 |
| $+$ LPU | 27M | 4.5G | 1065 | 83.3 |
| $+$ Re-Param | 27M | 4.5G | 1438 | 83.3 |
| $+$ MESA | 27M | 4.5G | 1438 | 83.5 |

Local enhancement. We replace the original non-overlapped stem and down-sampling layer with overlapped convolutions (refer to as ‘Overlapped Stem’ in Tab. 4), resulting in a gain of 0.2% in accuracy. We also incorporate the LPU (Guo et al., ), a depth-wise convolution placed at the top of each block and FFN, further increasing the accuracy by 0.4%. These modifications enrich the local information available between image patches before processing by the SASF module, enabling it to better capture structural dependencies.

Optimization. To further enhance the model efficiency, we implement the SASF module using re-parameterization techniques (Ding et al., ) and optimize the CUDA kernels. This accelerates the model by at least 30% and boosts the throughput from 1065 to 1438. Finally, integrating MESA (Du et al., ) to mitigate overfitting provides an additional 0.2% accuracy improvement.

In addition, we also provide some qualitative results in Appendix C, and a comparative analysis of the Effective Receptive Fields (ERF) (Ding et al., ) of various models is provided in Appendix D.

## 6 Conclusion

We presented in this paper Spatial-Mamba, a novel state space model designed for visual tasks. The key of Spatial-Mamba lied in the proposed structure-aware state fusion (SASF) module, which effectively captured image spatial dependencies and hence improved the contextual modeling capability. We performed extensive experiments on fundamental vision tasks of image classification, detection and segmentation. The results showed that with SASF, Spatial-Mamba surpassed the state-of-the-art state space models with only one signal scan, demonstrating its strong visual feature learning capability. We also analyzed in-depth the relationships of Spatial-Mamba with the original Mamba and linear attention, and unified them under the same matrix multiplication framework, offering a deeper understanding of the self-attention mechanism for visual representation learning.

## Appendix A Detailed Experiment Settings

Network architecture. The detailed architecture of Spatial-Mamba models are outlined in Tab. 5. Following the common four-stage hierarchical framework(Liu et al., ; Han et al., ), we constructed the Spatial-Mamba models by stacking proposed Spatial-Mamba blocks at each stage. Specifically, an input image with resolution of $224\times 224$ is firstly processed by a stem layer, which consists of convolutions, Batch Normalization (BN) layers and GELU activation function. The kernel size is $3\times 3$ with a stride of 2 at first and last convolutions and a stride of 1 for the others. Each stage contains multiple Spatial-Mamba blocks, following by a down-sampling layer except the last one. The down-sampling layer consists of a $3\times 3$ convolution with a stride of 2 and a Layer Normalization (LN) layer. Each block incorporates a structure-aware SSM layer and a Feed-Forward Network (FFN), both with residual connections. The structure-aware SSM contains a SASF branch with a 2D depth-wise convolution and a multiplicative gate branch with activation function, as illustrated in Fig. 4. The expand ratio of SSM is set to 2, doubling the number of channels. We modify the embedding dimension and number of blocks to build our Spatial-Mamba-T/S/B models.

Settings for ImageNet-1K classification. The Spatial-Mamba-T/S/B models are trained from scratch for 300 epochs using AdamW optimizer with betas set to (0.9, 0.999), momentum of 0.9 and batch size of 1024. The initial learning rate is set to 0.001 with a weight decay of 0.05 and a cosine annealing learning rate schedule is adopted with a warm-up of 20 epochs. We adopt the common data augmentation strategies consisted with previous work(Liu et al., ; ).
Moreover, label smoothing (0.1), exponential moving average (EMA) and MESA(Du et al., ) are also applied. The drop path rate is set to 0.2 for Spatial-Mamba-T, 0.3 for Spatial-Mamba-S and 0.5 for Spatial-Mamba-B.

**Table 5: Architectures of Spatial-Mamba models, where Linear refers to a linear layer, DWConv represents a depth-wise convolution layer, SASF module refers to Fig. 4 in Sec. 4.2 of the main paper, and FFN is a Feed-Forward Network.**
| Layer | Output size | Spatial-Mamba-T | Spatial-Mamba-S | Spatial-Mamba-B |
| --- | --- | --- | --- | --- |
| Stem | $56\times 56$ | Conv $3\times 3$ stride 2, BN, GELU; Conv $3\times 3$ stride 1, BN; Conv $3\times 3$ stride 2, BN |  |  |
| Stage1 | $28\times 28$ | Spatial-Mamba Blocks | Spatial-Mamba Blocks | Spatial-Mamba Blocks |
| $\left[\begin{array}[]{c}\text{Linear 64 $\rightarrow$ 128}\\<br>\text{DWConv 128}\\<br>\text{SASF 128}\\<br>\text{Linear 128 $\rightarrow$ 64}\\<br>\text{FFN 64}\end{array}\right]\times$ 2 | $\left[\begin{array}[]{c}\text{Linear 64 $\rightarrow$ 128}\\<br>\text{DWConv 128}\\<br>\text{SASF 128}\\<br>\text{Linear 128 $\rightarrow$ 64}\\<br>\text{FFN 64}\end{array}\right]\times$ 2 | $\left[\begin{array}[]{c}\text{Linear 96 $\rightarrow$ 192}\\<br>\text{DWConv 192}\\<br>\text{SASF 192}\\<br>\text{Linear 192 $\rightarrow$ 96}\\<br>\text{FFN 96}\end{array}\right]\times$ 2 |  |  |
| Down Sampling Conv $3\times 3$ stride 2, LN |  |  |  |  |
| Stage2 | $14\times 14$ | Spatial-Mamba Blocks | Spatial-Mamba Blocks | Spatial-Mamba Blocks |
| $\left[\begin{array}[]{c}\text{Linear 128 $\rightarrow$ 256}\\<br>\text{DWConv 256}\\<br>\text{SASF 256}\\<br>\text{Linear 256 $\rightarrow$ 128}\\<br>\text{FFN 128}\end{array}\right]\times$ 4 | $\left[\begin{array}[]{c}\text{Linear 128 $\rightarrow$ 256}\\<br>\text{DWConv 256}\\<br>\text{SASF 256}\\<br>\text{Linear 256 $\rightarrow$ 128}\\<br>\text{FFN 128}\end{array}\right]\times$ 4 | $\left[\begin{array}[]{c}\text{Linear 192 $\rightarrow$ 384}\\<br>\text{DWConv 384}\\<br>\text{SASF 384}\\<br>\text{Linear 384 $\rightarrow$ 192}\\<br>\text{FFN 192}\end{array}\right]\times$ 4 |  |  |
| Down Sampling Conv $3\times 3$ stride 2, LN |  |  |  |  |
| Stage3 | $7\times 7$ | Spatial-Mamba Blocks | Spatial-Mamba Blocks | Spatial-Mamba Blocks |
| $\left[\begin{array}[]{c}\text{Linear 256 $\rightarrow$ 512}\\<br>\text{DWConv 512}\\<br>\text{SASF 512}\\<br>\text{Linear 512 $\rightarrow$ 256}\\<br>\text{FFN 256}\end{array}\right]\times$ 8 | $\left[\begin{array}[]{c}\text{Linear 256 $\rightarrow$ 512}\\<br>\text{DWConv 512}\\<br>\text{SASF 512}\\<br>\text{Linear 512 $\rightarrow$ 256}\\<br>\text{FFN 256}\end{array}\right]\times$ 21 | $\left[\begin{array}[]{c}\text{Linear 384 $\rightarrow$ 768}\\<br>\text{DWConv 768}\\<br>\text{SASF 768}\\<br>\text{Linear 768 $\rightarrow$ 384}\\<br>\text{FFN 384}\end{array}\right]\times$ 21 |  |  |
| Down Sampling Conv $3\times 3$ stride 2, LN |  |  |  |  |
| Stage4 | $7\times 7$ | Spatial-Mamba Blocks | Spatial-Mamba Blocks | Spatial-Mamba Blocks |
| $\left[\begin{array}[]{c}\text{Linear 512 $\rightarrow$ 1024}\\<br>\text{DWConv 1024}\\<br>\text{SASF 1024}\\<br>\text{Linear 1024 $\rightarrow$ 512}\\<br>\text{FFN 512}\end{array}\right]\times$ 4 | $\left[\begin{array}[]{c}\text{Linear 512 $\rightarrow$ 1024}\\<br>\text{DWConv 1024}\\<br>\text{SASF 1024}\\<br>\text{Linear 1024 $\rightarrow$ 512}\\<br>\text{FFN 512}\end{array}\right]\times$ 5 | $\left[\begin{array}[]{c}\text{Linear 768 $\rightarrow$ 1536}\\<br>\text{DWConv 1536}\\<br>\text{SASF 1536}\\<br>\text{Linear 1536 $\rightarrow$ 768}\\<br>\text{FFN 768}\end{array}\right]\times$ 5 |  |  |
| Head | $1\times 1$ | Average pool, Linear 1000, Softmax |  |  |

## Appendix B Derivation of SSM Formulas

Recall the definition of Mamba in Eq. (1), by defining the initial state $x_{0}$ as zero, the state transition equation in the recursive form can be rewritten as follows:

$$ $\displaystyle x_{t}$ $\displaystyle=\overline{{\bm{A}}}_{t}x_{t-1}+\overline{{\bm{B}}}_{t}u_{t}$ (7) $\displaystyle=\overline{{\bm{A}}}_{t}\left(\overline{{\bm{A}}}_{t-1}x_{t-2}+ \overline{{\bm{B}}}_{t-1}u_{t-1}\right)+\overline{{\bm{B}}}_{t}u_{t}$ $\displaystyle=\overline{{\bm{A}}}_{t}\left(\overline{{\bm{A}}}_{t-1}\left( \overline{{\bm{A}}}_{t-2}x_{t-3}+\overline{{\bm{B}}}_{t-2}u_{t-2}\right)+ \overline{{\bm{B}}}_{t-1}u_{t-1}\right)+\overline{{\bm{B}}}_{t}u_{t}$ $\displaystyle=\overline{{\bm{A}}}_{t}\overline{{\bm{A}}}_{t-1}\overline{{\bm{A }}}_{t-2}x_{t-3}+\overline{{\bm{A}}}_{t}\overline{{\bm{A}}}_{t-1}\overline{{ \bm{B}}}_{t-2}u_{t-2}+\overline{{\bm{A}}}_{t}\overline{{\bm{B}}}_{t-1}u_{t-1}+ \overline{{\bm{B}}}_{t}u_{t}$ $\displaystyle=\overline{{\bm{A}}}_{t}\dots\overline{{\bm{A}}}_{1}x_{0}+ \overline{{\bm{A}}}_{t}\dots\overline{{\bm{A}}}_{2}\overline{{\bm{B}}}_{1}u_{1 }+\cdots+\overline{{\bm{A}}}_{t}\overline{{\bm{A}}}_{t-1}\overline{{\bm{B}}}_{ t-2}u_{t-2}+\overline{{\bm{A}}}_{t}\overline{{\bm{B}}}_{t-1}u_{t-1}+\overline{ {\bm{B}}}_{t}u_{t}$ $\displaystyle=\sum_{s\leq t}\overline{{\bm{A}}}_{s:t}^{\times}\ \overline{{\bm {B}}}_{s}u_{s},\quad\text{where }\overline{{\bm{A}}}_{s:t}^{\times}:=\begin{ cases}\overline{{\bm{A}}}_{t}\overline{{\bm{A}}}_{t-1}\dots\overline{{\bm{A}}} _{s+1}&s<t\\ 1&s=t\end{cases}.$ $$

By omitting term ${\bm{D}}_{t}u_{t}$ for simplicity, and multiplying by ${\bm{C}}_{t}$ according to the observation equation, we derive the final output $y_{t}=\sum_{s\leq t}{\bm{C}}_{t}\ \overline{{\bm{A}}}_{s:t}^{\times}\
\overline{{\bm{B}}}_{s}u_{s}$. Defining the input vector $u=\left[u_{1},\cdots,u_{L}\right]^{T}\in\mathbb{R}^{L}$ with sequences length $L$ and corresponding output vector $y\in\mathbb{R}^{L}$, then the above calculation can be written with in matrix multiplication framework, i.e. $y={\bm{M}}u$, where ${\bm{M}}$ is a structured lower triangular matrix and ${\bm{M}}_{ij}={\bm{C}}_{i}\overline{{\bm{A}}}_{j:i}^{\times}\overline{{\bm{B}}
}_{j}$.

Similarly, we can also represent Spatial-Mamba in the same matrix transformation form. Based on Eq. (2) and Eq. (7), the SASF equation can be rewritten as $h_{t}=\sum_{k\in\Omega}\sum_{s\leq\rho_{k}(t)}\alpha_{k}\overline{{\bm{A}}}_{s
:\rho_{k}(t)}^{\times}\overline{{\bm{B}}}_{s}u_{s}$. Then by multiplying ${\bm{C}}_{t}$ with out term ${\bm{D}}_{t}u_{t}$, we derive the final output of Spatial-Mamba, i.e. $y_{t}=\sum_{k\in\Omega}\sum_{s\leq\rho_{k}(t)}\alpha_{k}{\bm{C}}_{t}\overline{
{\bm{A}}}_{s:\rho_{k}(t)}^{\times}\overline{{\bm{B}}}_{s}u_{s}$. This can also be concisely represented as a matrix multiplication form $y={\bm{M}}u$, where ${\bm{M}}$ is a structured adjacency matrix and ${\bm{M}}_{ij}=\sum_{k}{\bm{C}}_{i}\overline{{\bm{A}}}_{j:\rho_{k}(i)}^{\times}
\overline{{\bm{B}}}_{j}$.

## Appendix C Effective Receptive Field (ERF)

We compare the Effective Receptive Field (ERF) of the center pixel on popular backbone networks before and after training, as shown in Fig. 6. Our Spatial-Mamba-T and Deit-S, Vim-S and VMamba-T all show a global ERF. In addition, after training, both Vim-S and VMamba-T exhibit noticeable accumulation contributions along either horizontal or vertical directions, which is attributed to their multi-directional fusion mechanisms. In contrast, our unidirectional Spatial-Mamba-T effectively eliminates this phenomenon.

Figure: Figure 6: Comparison of Effective Receptive Field (ERF) among popular backbone networks.
Refer to caption: x17.png

## Appendix D More Qualitative Results

In this section, we present the visualization results for object detection and instance segmentation tasks in Fig. 7 and for semantic segmentation task in Fig. 8. Compared with VMamba, our Spatial-Mamba demonstrates superior performance in both tasks, producing more accurate detections with higher confidence levels and more refined segmentations, particularly in areas where local structure information is crucial. For example, in the second row of Fig. 7, VMamba mistakenly identifies the shoes on a skateboard as a person, likely because observing the shoes from only four directions makes them resemble a human. Our method avoids this mistake by simultaneously perceiving the shoes and their surrounding context. Similarly, in the semantic segmentation task, as shown in the second and third rows of Fig. 8, our approach achieves more precise segmentations of the structure of trees and doors. These results highlight the effectiveness of our proposed Spatial-Mamba in leveraging local structural information for better visual understanding.

Figure: (a) Input
Refer to caption: x18.png

Figure: (a) Input
Refer to caption: x22.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 8
presentation: 4
questions: The results for Cascade Mask R-CNN seem to be missing...
Minor: On line 204, it is mentioned "considering both temporal and spatial information", but what exactly is the temporal information here? Or is it a typo?  

Overall, the paper is sound and presents a good idea for an existing problem. I am open to adjusting my evaluation should any significant concerns be raised by other reviewers that I may have overlooked.
soundness: 3
strengths: SSMs do have the problem of spatial sense. While the recent papers of ViM and VMamba try to solve them in their way, the solution feels rather underwhelming. The idea of introducing dilated convolutions as a means of capturing structural dependency in images is quite intriguing. I like the idea and the authors have done a good job of explaining most of the paper and related works. It is written clearly and concisely.
summary: Spatial-Mamba introduces a "structure-aware state fusion" (SASF) mechanism using dilated convolutions in the Mamba architecture. This mechanism claims to add neighborhood connectivity directly in the state-space of SSM. The method is benchmarked on common image benchmarks, i.e. segmentation, classification, and detection using common datasets.
weaknesses: There are not many weaknesses in this paper. Apart from Figure 5 (and qualitative and erf figures in supp.), the paper could improve on convincing the reader with visuals if the resultant SSM with SASF mechanism extracts structural dependencies from the images. 
Minor mistakes exist in the paper; specifically, on lines 199 and 200, equation 2 lacks a definition for alpha_k. It is only on line 212 that the alpha_k is mentioned as weights. Also, it is not clear from the text how exactly is this linear weighting being applied/calculated (on line 99, it mentions dilated convolutions to re-weight and merge states but it will be nice to add some clarity after equation 2 is introduced).

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: Please see weakness.
soundness: 3
strengths: 1. The motivation is reasonbal and good. The vision mamba models suffers from processing  2D visual data due to the complexity of spatial structures. It heavily relys  on the 1D scanning strategy to transform 2D format into 1D sequence. The author apply conv-based layers to fuse the contextual information  in spatial domain, which is complemenary to the Mamba based operator. 

2. The experimetanl results are good. The accuracy in imageNet shows the improvement. The speed is still competitive as the conv operators are optimized, thus bringing less latency. 

3. The overall method is simple and easy to understand. Applying conv layers as a compelmentary module to improve the spatial modeling capability of the mamba is simple and direct.
summary: This paper introduces Spatial-Mamba, a  state space models visual backbone designed for 2D visual tasks. It mainly introduce a multi-scale depthwise conv layer after the mamba operators to fuse the contextual information from different spatial locations. The author conducts theoretical analysis. Extensive experiements like image classification, dense prediction tasks are conducted.  The results show it outperforms the baselines and achieve competitive resutls.
weaknesses: 1. The method seems to be a combination of Mamba and a classical module (depthwise conv).  It is also mentioned in No.3 point in Strengths. I would recommned author to validate/explore the spatial modeling enhancement methods more throughly. For example, applying the multi-scale conv layer to the original bi-directional Mamba (Vim) and cross-scan mamba (VMamba) to see the improvements. THough the one-direction scan and conv layer can aggregate enough spatial contextual information. 

2. The analysis of linear attention, mamba and the proposed method in section 4.3 is good. But the method proposed is rather straightforward. It seems the method doesnot improve the weakness of Mamba from the in-depth analysis. It first uses the Selective scan in Mamba to calculate the hidden states (one direction), following a depthwise conv to fuse the hidden states, then calculate the output y. It is not known the improvements over the method that adding more direcitons into the selective scan and computing the output. 

3. The modules/operators to fuse the contextual information can have mutiple choices. The paper chooses depthwise conv. Conv operator may have inductive bias, which aggregate the spatial information by a fixed pattern. I would see more explorations, which could improve the paper.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 4
questions: Please refer to weakness part where I have expressed my concerns.
soundness: 4
strengths: + The proposed fusing operation seems to have a decent performance gain onto the baseline.
+ The manuscript is well-organized.
+ The ablation studies are clear and convincing.
summary: This manuscript proposes to aggregate neighborhood features in state-space models for more efficient state-space modeling. The experiments are convincing, and the manuscript is well-organized.
weaknesses: +  Novelty. The proposed fusion seems to be a re-invention and re-introduce of convolution on the feature map of state-space modeling but is claimed as "Structure-Aware State Fusion". The LPU, overlapped stem,  re-parameterization  and MESA are all proposed in previous works. The technical contribution seems to be limited.
+ Clearness. The fusion seems to be applied after scanning. However, in Eq(2), the $x_t$ is the same in scanning and fusion. The notation here is not clear. It would be better to use $x_t^{\prime}$ in the fusion formula if my understanding is correct.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. Could you further explain how state variables are modeled in the state space to design the structure-aware state fusion?
2. Please provide details on the number of hidden states in each SSM layer and include relevant ablation experiments.
3. Is Spatial-Mamba implemented using the hardware scanning algorithm proposed by Mamba?
soundness: 3
strengths: 1. Improved upon existing SSMs by directly integrating spatial neighborhood dependencies into the state transition process.

2. Analyzed the connections between Spatial-Mamba, Mamba, and linear attention, unifying them under a matrix multiplication framework.
summary: This paper introduces Spatial-Mamba, a novel state space model (SSM) designed for visual tasks, capable of effectively capturing spatial dependencies in images to enhance contextual modeling. Experimental results demonstrate that Spatial-Mamba outperforms existing SSMs and other mainstream models in areas such as image classification and semantic segmentation.
weaknesses: 1. The paper does not mention the number of hidden states used in the SSM, and the ablation study lacks an analysis of the impact of different numbers of hidden states.

2. While the authors analyze the relationship between Spatial-Mamba, Mamba, and linear attention, they do not clarify the framework used to implement Spatial-Mamba, such as Mamba's hardware scanning algorithm or Mamba2's SSD framework.

3. In Mamba’s proposed SSM, hidden states do not inherently have spatial relationships. I suspect the authors intend to model the hidden states as 2D local observation points?

Meta-review: This paper introduces Spatial-Mamba, a backbone based on State Space Modeling (SSM) specifically designed for 2D vision tasks. It establishes neighborhood connectivity within the state space through a method known as structure-aware state fusion. 

The strengths of the paper include a clear motivation, solid experimental results, and a straightforward method that is easy to understand. However, there are some concerns regarding the novelty of the work and the lack of comprehensive ablation studies.

The paper received generally positive reviews with scores of 8, 8, 6, and 6, resulting in an average score of 7. Since all reviewers acknowledged the paper's strengths, the Area Chair recommends its acceptance.

Decision: Accept (Poster)