SF-Mamba: Rethinking State Space Model for Vision
Abstract
Visual Mamba models have recently emerged as alternatives to Vision Transformers (ViTs), which suffer from quadratic complexity. While the recurrent scanning mechanism of Mamba offers computational efficiency, it inherently limits non-causal interactions between image patches. Prior works have attempted to address this limitation through various multi-scan strategies; however, these approaches suffer from inefficiencies due to suboptimal scan designs and frequent data rearrangement. Moreover, Mamba exhibits relatively slow computational speed under short token lengths, commonly used in visual tasks. In pursuit of a truly efficient vision encoder, we rethink the scan operation for vision and the computational efficiency of Mamba. To this end, we propose SF-Mamba, a novel visual Mamba with two key proposals: auxiliary patch swapping for encoding bidirectional information flow under an unidirectional scan and batch folding with periodic state reset for advanced GPU parallelism. Extensive experiments on image classification, object detection, and instance and semantic segmentation consistently demonstrate that our proposed SF-Mamba significantly outperforms state-of-the-art baselines while improving throughput across different model sizes. The source code is available at: https://github.com/sony/SF-Mamba.
Lay Summary
We found that existing Visual Mamba models suffer from slow inference speed, particularly on low-resolution images. Through investigating the cause, we identified two major sources of inefficiency in current Visual Mamba architectures. First, prior works rely on multi-directional scans to alleviate the causality constraint of Mamba. However, we found that the associated data rearrangement operations introduce unexpectedly large overheads. Second, Mamba itself becomes inefficient under short token lengths, which are common in low-resolution visual tasks. To address the first issue, we propose auxiliary token swapping, which enables global information propagation while preserving a uni-directional scan scheme. To address the second issue, we introduce batch folding with periodic state reset, which improves the efficiency of Mamba under short token lengths. As a result, our method achieves a superior accuracy-throughput trade-off compared to existing approaches across both low- and high-resolution settings.