Source-linked AI summary
LocalMamba: Visual State Space Model with Windowed Selective Scan
Tao Huang, Xiaohuan Pei, Shan You, Fei Wang, Chen Qian, Chang Xu
TL;DR
Vision Mamba models have not clearly surpassed CNNs and ViTs, partly because flattening images disrupts local 2D dependencies. LocalMamba uses windowed selective scanning and layer-wise direction search to combine local detail with global context, and reports improvements across vision tasks, including a 1.5 mIoU (SS) margin over Vim-S on semantic segmentation.
Problem
Flattening 2D images into 1D sequences disrupts local spatial dependencies, while vision SSMs have shown only modest improvements over CNNs and ViTs.
Method
LocalMamba partitions images into windows, combines local and global selective scans with SCAttn, and searches scan directions independently across layers.
Results
Across image classification, object detection, and semantic segmentation, LocalMamba improves over previous methods; LocalVim-S exceeds Vim-S by 1.5 mIoU (SS) at similar parameter counts.
Takeaways & Limitations
The results support windowed scanning and layer-specific direction selection as effective design components for visual state space models.
Takeaways & Limitations
The direction-search approach scales poorly for extensive choices, requiring 28 GB of GPU memory for 20 blocks with 128 directions per block.
Abstract
from arXiv · showhide
Recent advancements in state space models, notably Mamba, have demonstrated significant progress in modeling long sequences for tasks like language understanding. Yet, their application in vision tasks has not markedly surpassed the performance of traditional Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs). This paper posits that the key to enhancing Vision Mamba (ViM) lies in optimizing scan directions for sequence modeling. Traditional ViM approaches, which flatten spatial tokens, overlook the preservation of local 2D dependencies, thereby elongating the distance between adjacent tokens. We introduce a novel local scanning strategy that divides images into distinct windows, effectively capturing local dependencies while maintaining a global perspective. Additionally, acknowledging the varying preferences for scan patterns across different network layers, we propose a dynamic method to independently search for the optimal scan choices for each layer, substantially improving performance. Extensive experiments across both plain and hierarchical models underscore our approach's superiority in effectively capturing image representations. For example, our model significantly outperforms Vim-Ti by 3.1% on ImageNet with the same 1.5G FLOPs. Code is available at: https://github.com/hunto/LocalMamba.
1 Introduction
Vision SSMs have struggled to preserve local 2D image dependencies after flattening images into sequences. LocalMamba addresses this with windowed scanning, direction search, and architectures that combine local and global information.
- Motivation: Vision SSMs retain efficient sequence processing but have shown only modest gains over CNNs and ViTs on vision tasks.Flattening 2D images into 1D tokens disrupts local spatial dependencies and weakens spatial relationship modeling.
- Local scanning: Windowed scanning partitions images into local regions so semantically related tokens are processed closely together.The method then traverses across windows, improving local representation while preserving a global perspective.
- Architecture: LocalMamba combines global scanning directions with local scanning to capture comprehensive context and fine-grained regional information.Its SCAttn module adaptively emphasizes useful spatial and channel features while filtering redundancy.
- Adaptive directions: A learnable direction-search method selects effective scan patterns independently across network layers.The motivation is that scan preferences vary with layer depth and object scale, such as smaller windows for details and larger windows for larger objects.
- Results: 1.5 mIoU (SS) is the reported margin by which LocalVim-S outperforms Vim-S with a similar parameter count on semantic segmentation.The paper evaluates plain and hierarchical variants across image classification, object detection, and semantic segmentation.
2 Related Work
The related work traces computer vision from CNNs and ViTs to state space models, emphasizing SSMs' sequence modeling efficiency and growing use in visual tasks.
- CNNs and ViTs: CNNs established major visual architectures, while ResNet and MobileNet emphasized residual learning and lightweight depth-wise convolutions.These models became widely used across vision tasks before transformer-based approaches gained prominence.
- CNNs and ViTs: Vision Transformers process images as patch sequences using self-attention, challenging the earlier dominance of CNNs.The passage characterizes ViT as a major shift in image-processing architecture.
- State space models: SSMs map sequences while targeting long dependencies, and later advances improved their computational and memory efficiency.S4 introduced an efficient Normal Plus Low-Rank representation to streamline computation.
- Visual SSMs: Visual SSM research progressed from S4ND to Vim and VMamba, with later models using bidirectional or cross-scan mechanisms for visual data.S4ND addressed continuous visual signals across 1D, 2D, and 3D domains, while Vim and VMamba targeted generic vision tasks.
3 Preliminaries
The preliminaries describe SSMs as sequence-to-sequence systems with latent states, then introduce selective SSMs, or Mamba, whose parameters depend on the input sequence.
- SSM formulation: An SSM maps a one-dimensional input sequence to an output sequence through an intermediate latent state.Its system matrices govern state dynamics and output mapping.
- Discretization: Continuous SSM dynamics are discretized with a zero-order hold assumption before practical recurrent computation.The discretized formulation updates the latent state from the previous state and current input, then produces the output.
- Efficient computation: SSMs can compute sequence outputs efficiently through a global convolution using an SSM kernel.The convolution synthesizes outputs simultaneously, supporting computational efficiency and scalability.
- Selective SSMs: Traditional SSMs use static parameterization, limiting context capture despite linear-time complexity.This limitation motivates selective state space models.
- Selective SSMs: Mamba uses input-dependent B, C, and Δ parameters to make state interactions dynamic and sequence-aware.These parameters are calculated directly from the input sequence.
- Visual applications: Vim and VMamba extend Mamba to vision through bidirectional blocks and horizontal-vertical 2D selective scanning.These approaches integrate selective SSMs into generic vision architectures.
4 Methodology
LocalMamba rearranges image tokens into windows to improve local dependency modeling, combines local and global scans with attention, and searches scan directions per layer.
- 4.1 Local Scan for Visual Representations: Local scanning divides images into windows, placing relevant local tokens closer together for selective processing.The design targets the loss of spatial coherence caused by flattening images into 1D sequences.
- 4.1 Local Scan for Visual Representations: The LocalMamba block combines original and flipped scan directions across four selective-scan branches to retain local and global information.Branch outputs are merged after independent feature extraction and attention-based aggregation.
- 4.1 Local Scan for Visual Representations: Unlike windowed self-attention in ViTs, LocalMamba rearranges token positions for local dependencies while the full image remains aggregated by the SSM.The distinction is therefore about dependency modeling rather than reducing global attention computation.
- Scalability: Direction search has scalability limits: a 20-block model with 128 directions per block requires 28 GB of GPU memory.The authors identify sampling, binary approximation, and partial-channel methods as possible mitigations and leave more adaptive strategies for future work.
- Architecture variants: LocalVim and LocalVMamba adapt the approach to plain and hierarchical architectures, respectively.LocalVim replaces the standard SSM block and adjusts block count for computation, while LocalVMamba preserves its structural configuration.
- Computational cost: The scan itself adds no FLOPs because it only repositions tokens, while SCAttn introduces a streamlined aggregation module.The supplied passage reports only a marginal computation increase for the LocalMamba block.
5 Experiments
Experiments evaluate LocalVim and LocalVMamba on ImageNet classification, COCO detection and instance segmentation, and ADE20K semantic segmentation. Ablations show that local scans, scan-direction search, and SCAttn each contribute to performance gains.
- ImageNet Classification: 76.2% ImageNet accuracy at 1.5G FLOPs lets LocalVim-T surpass DeiT-Ti’s 72.2%.
- ImageNet Classification: 82.7% accuracy makes LocalVMamba-T outperform Swin-T by 1.4% on ImageNet classification.
- Object Detection: LocalVMamba-T reaches 46.7 box AP and 42.2 mask AP on COCO, improving Swin-T by 4.0 and 2.9, respectively.
- Semantic Segmentation: LocalVim-S exceeds Vim-S by 1.5 mIoU (SS), while LocalVMamba-T reaches 49.1 mIoU (MS), surpassing VMamba-T by 0.8.
- Ablation Study: Replacing Vim-T’s horizontal scan with a local scan boosts performance by 1%, and combining scan directions adds another 1.1% under a constrained FLOP budget.
- Ablation Study: SCAttn adds 0.6%, while scan-direction search improves LocalVim-T over LocalVim-T∗ by 0.4%.
- Visualization of Searched Scan Directions: Searched directions vary by architecture and depth: plain LocalVim favors larger scans near the beginning and smaller scans near the tail.
6 Conclusion
LocalMamba uses windowed selective scanning and scan-direction search to capture local dependencies while preserving global context. Experiments across datasets and tasks report superiority over traditional CNNs and ViTs.
- LocalMamba enhances local dependency capture while maintaining global contextual understanding through windowed selective scanning and scan-direction search.
- The approach is evaluated across various datasets and tasks, including image classification, object detection, and semantic segmentation.
- The authors report that LocalMamba establishes new benchmarks and outperforms traditional CNNs and ViTs.
A.1 Comparison to Vim on Object Detection
LocalVim is evaluated against Vim under aligned detection settings on the COCO validation set. It performs comparably overall while improving selected box and mask metrics.
- The evaluation uses Vim-aligned settings, with Vim employing ViTDet’s neck and Cascade Mask R-CNN rather than VMamba’s Mask R-CNN framework.
- LocalVim-T performs on par with Vim-Ti while showing significant advantages in APb50 and mask AP.
- LocalVim-T improves Vim-Ti on APm and APm50 by 0.7 and 2.1, respectively.
A.2 Visualization of Searched Directions on LocalVMamba-S
The searched scan directions in LocalVMamba-S vary with network depth. In stage 3, the deeper model favors more large-window local scans than LocalVMamba-T.
- Figure 5 visualizes the searched directions of LocalVMamba-S.
- LocalVMamba-S contains 27 layers in stage 3.
- Compared with LocalVMamba-T, LocalVMamba-S prefers more 7 × 7 local scans.
A.3 Discussions
The discussion identifies computational efficiency as an important boundary for visual state space models. Although SSMs scale linearly with sequence length, their execution remains harder to accelerate efficiently than convolution and self-attention.
- The proposed model requires substantial computational resources, which may raise environmental concerns.
- SSMs have linear-time complexity with sequence length and show especially significant improvements on large-resolution downstream tasks.
- SSM computation is more intricate than convolution and self-attention, complicating parallel execution and limiting current framework acceleration.