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

We sincerely thank you for your constructive comments! Please find our responses to your concerns below.

---

Q1. Analysis on the number of hidden states.

Thank you for your insightful comments. We mainly follow the configurations of VMamba for fair comparison with the competing methods. The number of hidden states (the SSM state dimension, i.e., d_state) is set to 1. The impact of the number of d_state has been discussed in VMamba. It is found that gradually increasing d_state from 1 to 4 slightly improves the performance, but the throughput drops significantly, indicating a negative impact on computational efficiency. We also analyzed the setting of d_state in our early attempts (single directional scanning without SASF and MLP), and the results are shown in the Table below. We see that increasing d_state from 16 to 32 improves the performance by 0.3% in top-1 accuracy on ImageNet-1K classification, but the throughput drops by 30%. Therefore, in all our experiments, we set d_state to 1 as the default value.

Table4. Ablation study of the number of hidden states d_state.

Mamba2 solves the above problem, which allows the use of a larger d_state (64 or 128) while maintaining efficient throughput. Nonetheless, this is not our focus. In future work, we plan to extend our study by incorporating more ablation experiments on the number of hidden states under the Mamba2 architecture.

---

Q2. Clarification on the framework of SSM and implementation details.

Thank you for your valuable comments. We have provided more details about the framework used in Spatial-Mamba in the revised version. We employ a hardware-aware selective scan algorithm, which is based on the original Mamba framework with some modifications. Specifically, we made modifications in the CUDA kernels to decouple the state transition and the observation equations. Besides, the SASF module of Spatial-Mamba is implemented by the general matrix multiplication (GEMM) with optimized CUDA kernels.

---

Q3. The spatial relationships of hidden states.

As mentioned by this reviewer, the hidden states in the original Mamba, which was designed for 1D sequential data, cannot well encode the spatial relationships that are crucial for understanding 2D images. In 1D sequences, the relationships between elements are primarily defined by their sequential order. However, in 2D images, the relationships between pixels or states are much more complex, involving both horizontal and vertical dependencies. To address this limitation and enhance the expressive power of the state space, we establish connections between neighboring states in the 2D state space. This explicit modeling of local spatial relationships allows the model to capture the inherent structure of the image data, leading to a more informative and powerful representation. By considering the interactions between nearby states, our proposed method effectively integrates the spatial context that is neglected in the original Mamba.

---

Q4. Explanation on how state variables are modeled in SASF.

The key challenge in extending Mamba from 1D data to 2D images is how to establish spatial relationships between state variables. Intuitively, we could attend each state variable to its neighboring states, enabling spatial structure awareness in the state space. To achieve this, we introduce structure-aware state fusion equation between the state transition and observation equations (see Fig. 1(d) and Eq. (2) in the main paper). Unlike the original Mamba, where state variables are only connected to the previous state, SASF allows each state to capture information from neighboring states through dilated convolution, converting it to a structure-aware state variable. This SASF stage captures the spatial dependencies between neighboring states in 2D data. We will add more explanations in the revised manuscript.

We sincerely thank you for your constructive comments and recognition on our work! We hope our following responses can address your concerns.

---

Q1. Clarification on novelty.

Thanks for your comments. We are sorry if we did not explain our contributions clearly enough in the manuscript, and we would like to clarify the novelties of our work in this response.

The SASF module is not a re-introduction of convolution, but an innovative design on how we can effectively fuse states to capture complex spatial relationships while preserving the advantages of the original Mamba. As indicated by this reviewer, LPU, overlapped stem, re-parameterization and MESA have been explored in prior works, and we just adopt them in the implementation of our model to further improve its performance. Our main contribution lies in the visual state space modeling through SASF to capture spatial dependencies in 2D data. As we responded to the third question of reviewer Q4w4, SASF can be implemented by different operators beyond depth-wise convolution. In addition, we have analyzed the impact of overlapped stem, LPU, re-parameterization, and MESA in the ablation study. The results showed that that even without using these modules, Spatial-Mamba-T still outperforms the baseline by 0.7% and VMamba-T (with overlapped stem) by 0.3% in top-1 accuracy on ImageNet-1K. This validates the effectiveness of SASF.

---

Q2. Clarification on Eq. (2).

Thanks for your comments. The SASF is applied after the scanning phase but before the observation equation in the model. We would like to clarify that the notation used in Eq. (2) is accurate and appropriate. For your convenience, we briefly summarize the processes of Mamba and Spatial-Mamba in this response. The standard Mamba workflow processes input $u_t$ through the state transition phase (the scanning process) to generate $x_t$, which is then passed directly into the observation phase to produce the output $y_t$ $\mathbf{(}u_t \to x_t \to y_t \mathbf{)}$. In contrast, our Spatial-Mamba workflow introduces an additional step: input $u_t$ is processed through the scanning phase to generate $x_t$, which then goes through the SASF phase to produce the fused state $h_t$. Finally, $h_t$ is passed into the observation equation to yield $y_t$ $\mathbf{(}u_t \to x_t \to {\color{red} h_t} \to y_t \mathbf{)}$. Therefore, the notation used in Eq. (2) is accurate and appropriate.

We sincerely thank you for your constructive comments! Please find our response to your concerns below.

---

Q1. Further exploration of spatial modeling ability.

Thanks for your constructive comments. We'd like to clarify that our proposed Spatial Mamba is not a simple combination of Mamba and depth-wise conv, while depth-wise conv is an efficient way to implement our proposed SASF. Our main contribution lies in the innovative design of SASF, which considers both the global long-range dependencies and local spatial relationships within the state space. As mentioned in our manuscript, while existing multi-directional models (like those using bi-directional and cross-scan models) possess certain spatial modeling capability, they introduce much more additional computational costs, limiting their overall performance. Our Spatial-Mamba, with a single scan, achieves both improved accuracy and faster inference than existing visual Mamba models.

Based on this reviewer's suggestion, to further explore and validate SASF's spatial modeling capability, we adapt SASF into Vim and VMamba backbones, resulting in Vim+SASF (by replacing the middle class token with average pooling and setting state dimension to 1) and VMamba+SASF (directly integrated with SASF). Under the same training settings, the classification results on ImageNet-1K are presented in the following Table:

Table2. Comparison of classification performance on ImageNet-1K.

It can be seen that by integrating with our proposed SASF, the bi-directional Vim model is improved by 0.5% in top-1 accuracy and the cross-scan VMamba is improved by 0.3%. These findings verify that SASF can even enhance the performance of multi-directional models, which have already incorporated (yet in an inefficient manner) spatial modeling operations. We will supplement these experiments in the ablation study of the revised manuscript.

---

Q2. In-depth analysis of the proposed method.

Thanks for your comments. We'd like to clarify that our proposed Spatial-Mamba indeed overcomes the weakness of original Mamba, at least to certain extent, in processing 2D visual data by introducing the SASF equation, while we provide a simple yet and effective solution to implement SASF via depth-wise convolution. In Sec. 1 and Fig. 1, Sec. 4.1 and Fig. 3, and Sec. 4.3 and Fig. 5 of the manuscript, we have analyzed the weaknesses of Mamba in adapting to 2D data and how Spatial-Mamba can address them. In specific:

• Sec. 1 and Fig. 1 outline the limitations of Mamba’s directional scanning in preserving 2D spatial relationships and its computational inefficiencies.
• Sec. 4.1 and Fig. 3 demonstrate how Spatial-Mamba explicitly models spatial relationships within the state space, effectively overcoming Mamba’s limitations in processing 2D data.
• Sec. 4.3 and Fig. 5 present a detailed comparative analysis, showing how Spatial-Mamba achieves superior spatial representation by leveraging weighted summation over a broader neighborhood, as visualized through adjacency matrices.

To further clarify, we have updated Fig. 5 in the manuscript for better illustration. By rearranging the elements of certain rows of $\mathbf{M}$ (indicated by the red arrow) into image space, we illustrate the activation map of the corresponding patch (marked with a red star) on the right side of the figure. This visualization demonstrates that linear attention focuses on a limited area, while Mamba captures a broader region due to its long-range context modeling. Our Spatial-Mamba not only largely extends the range of context modeling but also enables spatial structure modeling, effectively identifying relevant regions even when they are distant from each other.

Thank you for your insightful comments! Please find our responses below.

---

Q1. More visual results.

We sincerely thank this reviewer for the recognition on our work. Actually, in Fig. 3 of our manuscript, we have provided a visualization on the differences in state variables before and after our SASF, highlighting SASF's effectiveness in spatial modeling. Based on the suggestion of this reviewer, we have provided more visual results in Fig. 1 of the updated supplementary materials. We can see that the fused state variables show more accurate context information and structural perception capabilities across different scenarios.

---

Q2. Definition of $\alpha_k$ and clarification on the linear weighting.

We are sorry for the missing description of $\alpha_k$ in Eq. (2). $\alpha_k$ represents the learnable weight used in our weighting method. As illustrated in Fig. 1(d) of the main paper, neighboring state variables (shown in blue and orange colors) are aggregated using these weights to produce a fused state variable (shown in red color). This fusion process is efficiently implemented using dilated convolutions. We will add these descriptions and make the definition of $\alpha_k$ and the calculation of linear weighting clearer in the revised manuscript.

---

Q3. Results of Cascade Mask R-CNN.

Thanks for the suggestion. We followed previous works (VMamba, LocalMamba) to use Mask R-CNN as the detector head in the experiments for a fair comparison with them. As suggested by this reviewer, we further performed experiments on Spatial-Mamba under the Cascade Mask R-CNN framework. We adopt the same configurations and hyper-parameter settings as Swin Transformer and fine-tune the pre-trained Spatial-Mamba-T/S models for 36 epochs (3x schedule) with Cascade Mask R-CNN detector head on the COCO dataset. The results are shown in the Table below:

Table1. Cascade Mask R-CNN - 3x schedule

We can see that Spatial-Mamba-T achieves a box mAP of 52.1 and a mask mAP of 44.9, surpassing Swin-T/NAT-T by 1.7/0.7 in box mAP and 1.2/0.4 in mask mAP with fewer parameters and FLOPs, respectively. Similarly, Spatial-Mamba-S demonstrates superior performance under the same configuration. We will add these results in the revision.

(Note$^*$: The results of Vim shown in the first row of the table are for reference only, because in the original paper of Vim, only the results of the Vim-tiny variant (7M) with Cascade Mask R-CNN are reported, and the training framework is also different from our 3x schedule.)

---

Q4. Explanation on the temporal information.

We are sorry for the confusion caused. As mentioned in the manuscript, compared with the original Mamba, we enhance its spatial modeling capabilities by introducing a structure-aware state fusion (SASF) equation after the state transition equation. Our intention of using "temporal" here is to express that we retain the advantages of the original Mamba in processing temporal data, such as its ability to capture long-range and selective information. We will revise the sentence "by considering both temporal and spatial information" to "by considering both the global long-range and the local spatial information" in order to avoid possible misunderstandings.

I appreciate the author’s response.

For Fig. 3, I am unsure how exactly in Fig. 3 (b) and Fig. 3 (c) the high mean of state variables in non-interesting regions can be interpreted as the method's effectiveness.

Thank you for including the 3x schedule results.

My other concerns, as well as those raised by other reviewers, have been mostly addressed. I will maintain my recommendation for accepting this paper.