Source-linked AI summary
PlainMamba: Improving Non-Hierarchical Mamba in Visual Recognition
Chenhongyi Yang, Zehui Chen, Miguel Espinosa, Linus Ericsson, Zhenyu Wang, Jiaming Liu, Elliot J. Crowley
TL;DR
PlainMamba addresses the challenge of adapting Mamba’s sequential selective scanning to two-dimensional visual recognition while preserving a simple, reusable architecture. It introduces continuous 2D scanning and direction-aware updating in a constant-width stack of identical blocks, and reports gains over non-hierarchical models with competitiveness against hierarchical alternatives. Its main boundary is weaker performance than hierarchical VMamba in the reported object-level setting.
Problem
Adapting Mamba’s selective scanning to images requires handling two-dimensional spatial relationships while retaining competitive visual representations and practical architectural simplicity.
Method
PlainMamba uses identical constant-width blocks, removes special tokens, and adapts selective scanning with continuous 2D scanning and direction-aware updating.
Results
PlainMamba outperforms previous non-hierarchical models across visual recognition tasks and performs competitively with hierarchical alternatives.
Takeaways & Limitations
PlainMamba provides a simple, scalable baseline for visual recognition, including high-resolution tasks where it maintains high performance with less computation.
Takeaways & Limitations
PlainMamba performs worse than hierarchical VMamba in the reported object-level setting, which the authors relate to the multi-resolution architecture used by FPN-based Mask R-CNN.
Abstract
from arXiv · showhide
We present PlainMamba: a simple non-hierarchical state space model (SSM) designed for general visual recognition. The recent Mamba model has shown how SSMs can be highly competitive with other architectures on sequential data and initial attempts have been made to apply it to images. In this paper, we further adapt the selective scanning process of Mamba to the visual domain, enhancing its ability to learn features from two-dimensional images by (i) a continuous 2D scanning process that improves spatial continuity by ensuring adjacency of tokens in the scanning sequence, and (ii) direction-aware updating which enables the model to discern the spatial relations of tokens by encoding directional information. Our architecture is designed to be easy to use and easy to scale, formed by stacking identical PlainMamba blocks, resulting in a model with constant width throughout all layers. The architecture is further simplified by removing the need for special tokens. We evaluate PlainMamba on a variety of visual recognition tasks, achieving performance gains over previous non-hierarchical models and is competitive with hierarchical alternatives. For tasks requiring high-resolution inputs, in particular, PlainMamba requires much less computing while maintaining high performance. Code and models are available at: https://github.com/ChenhongyiYang/PlainMamba .
1 Introduction
PlainMamba addresses the tension between benchmark performance and the simplicity and reusability of plain visual encoders by adapting Mamba for visual recognition. It uses a constant-width, non-hierarchical design with 2D-aware scanning and reports competitive performance across visual tasks.
- Motivation and contribution: Plain non-hierarchical encoders are widely used because their simple structure improves reuse, despite hierarchical encoders sometimes achieving higher benchmark accuracy.The paper argues that benchmark performance may not fully reflect real-world usefulness when additional engineering increases complexity.
- Background: SSMs offer potentially infinite context lengths with linear complexity in sequence length, motivating their adaptation from sequential data to visual recognition.The paper places Mamba within broader efforts to build competitive SSM representations.
- Motivation and contribution: PlainMamba is a simple visual SSM that maintains constant width by stacking identical blocks and removes the need for CLS tokens.The design targets easy scaling and model reuse while achieving competitive performance.
- Motivation and contribution: PlainMamba adapts selective scanning to 2D inputs through continuous spatially adjacent scanning and direction-aware updating.These techniques aim to preserve semantic continuity and encode the directionality of scanning orders.
- Evaluation: PlainMamba is evaluated in three sizes across ImageNet-1K classification, semantic segmentation, and object detection, outperforming non-hierarchical counterparts while matching hierarchical competitors.The reported comparisons span SSMs, transformers, and CNN-based visual feature extractors.
2 Related Work
Prior visual recognition research spans CNNs, vision transformers, SSMs, and simplified plain architectures. These lines of work trade spatial inductive biases, sequence-processing efficiency, and architectural simplicity in different ways.
- Visual feature extractors: CNNs exploit local visual features through convolutional spatial priors, while deeper stacks enlarge the receptive field.Their architecture has long dominated visual feature extraction.
- PlainMamba: PlainMamba follows this simplification trend with identical blocks, global average pooling instead of CLS, and direction-aware updating for relative 2D position.Its architecture combines selective scanning with a gated MLP.
- Simplifying visual feature extractors: Plain architectures simplify feature reuse by avoiding hierarchical pyramids, and plain ViTs provide flexibility for pre-training and fine-tuning.Related methods also unify input resolutions or improve feature resolution without introducing hierarchy.
- State space models: SSMs provide transformer alternatives by modeling extended sequence dependencies with linear scaling in sequence length.Prior work improves SSM efficiency through low-rank constraints, parallel scanning, and better hardware utilization.
- Visual feature extractors: Vision transformers replace CNN pyramid structures with patched tokens and learnable contexts, enabling flexible adaptation across visual tasks.Plain ViTs are also used as general-purpose image encoders in large-scale systems.
3 Method
PlainMamba adapts selective state-space scanning to 2D images while simplifying the backbone into identical, constant-width blocks without special tokens. Its main changes preserve spatial continuity and encode scan direction when processing flattened visual tokens.
- Preliminaries: Mamba makes B, C, and Δ depend on the input sequence, producing token-dependent parameters for selective scanning and combining the SSM with a gated MLP.The underlying SSM maps sequential inputs to outputs through recurrent hidden-state updates; σ is SiLU and ⊙ is element-wise multiplication.
- Architecture: PlainMamba maps images to visual tokens, adds positional embeddings, and processes them through L identical blocks with consistent input-output shapes.A convolutional tokenizer downsamples the input by the default factor of 16 before the block stack.
- Architecture: The backbone maintains constant feature resolution and model width across blocks, removes special tokens, and supports scaling by changing width and depth.Variants scale width first and then both width and depth; FLOPs are measured at 224×224 inputs.
- Continuous 2D Scanning: Continuous 2D Scanning uses four orders in which every token is spatially adjacent to its predecessor, including at row or column transitions.The scan reverses direction at transitions to preserve spatial and semantic continuity.
- Direction-Aware Updating: Direction-Aware Updating adds direction-specific information to selective scanning because each token may occupy one of four directions relative to its predecessor.The formulation uses direction-indexed parameters, with a separate final vector for the initial token of each scan.
4 Experiments
PlainMamba is evaluated across image classification, semantic segmentation, object detection, and instance segmentation against CNN, transformer, and SSM backbones. It performs strongly among non-hierarchical models and remains competitive with hierarchical alternatives, although hierarchical designs are better for some object-level evaluations.
- Experimental scope: PlainMamba is compared across four visual recognition tasks using CNN, vision transformer, and SSM feature extractors.The evaluation covers image classification, object detection, instance segmentation, and semantic segmentation.
- Evaluation settings: ImageNet-1K FLOPs are compared in Table 2, while ADE20K FLOPs use input size 512×2048 and Mask R-CNN FLOPs use 1280×800.The ADE20K setup uses UperNet; the COCO setup uses ViTAdapter to compute multi-scale features for FPN.
- ImageNet-1K Classification: 2.4% higher accuracy: PlainMamba-L2 outperforms Mamba-ND-T at similar model size on ImageNet-1K.PlainMamba-L2 also achieves 1.7% better accuracy than DeiT-Small and 0.3% better accuracy than Swin-Tiny around 25M parameters.
- ADE20K Semantic Segmentation: 1.9 higher mIoU: PlainMamba-L2 outperforms ViM-S on ADE20K semantic segmentation.PlainMamba-L2 also surpasses XCiT-S12/8 with lower computation and Swin-Tiny with lower model size and FLOPs.
- COCO Detection and Instance Segmentation: 44.1 AP𝑏𝑏 and 39.1 AP𝑚𝑘: PlainMamba-L1 trails Swin-Small while using similar FLOPs and many fewer parameters under the 1× Mask R-CNN schedule.PlainMamba nevertheless achieves 3.1 higher AP𝑏𝑏 and 1.6 higher AP𝑚𝑘 than DeiT-T with ViTAdapter.
- COCO Detection and Instance Segmentation: Hierarchical models tend to perform better than non-hierarchical models on the Mask R-CNN object-level evaluation.The authors attribute PlainMamba’s inferiority to VMamba to the multi-resolution FPN-based architecture being more suitable for hierarchical designs.
5 Conclusion
PlainMamba is presented as a simple plain SSM-based visual recognition model without special tokens or hierarchical structure. The reported results show stronger performance than previous non-hierarchical models and performance on par with high-performing hierarchical models.
- 5 Conclusion: PlainMamba uses no special tokens or hierarchical structure, providing a plain SSM-based counterpart to the widely used plain vision transformer.The architecture is intended as a baseline for future visual state space model research.
- 5 Conclusion: PlainMamba achieves superior performance to previous non-hierarchical models, including concurrent SSM-based models.The conclusion reports this pattern across the paper’s visual recognition evaluations.
- 5 Conclusion: PlainMamba performs on par with high-performing hierarchical models.The conclusion frames PlainMamba as a simple baseline despite using a non-hierarchical design.
A.1 COCO Object Detection using RetinaNet
PlainMamba is also evaluated with the single-stage RetinaNet detector on COCO. It achieves slightly higher AP than Swin-Tiny with half the model size and similar FLOPs.
- A.1 COCO Object Detection using RetinaNet: 0.2 higher AP: PlainMamba-L1 surpasses Swin-Tiny with only half the model size and similar FLOPs.The comparison uses the COCO RetinaNet object detection evaluation.
- A.1 COCO Object Detection using RetinaNet: RetinaNet results are reported on MS COCO mini-val using a 1× schedule, with FLOPs computed at input size 1280×800.The experiment tests PlainMamba in a single-stage object detector.
A.2 Ablation Studies and Discussions
The ablations examine PlainMamba’s depth–width balance, block design, and efficiency against DeiT. Deeper models improve accuracy up to a point, while the proposed block and high-resolution scanning offer favorable performance and computational trade-offs.
- Depth v.s. Width: At roughly 7.4M parameters, 12 layers improve ImageNet top-1 accuracy by 2.2% over 6 layers, and 24 layers add another 1.1% over 12 layers.
- Depth v.s. Width: The 36-layer model maintains similar top-1 accuracy after reducing width, while deeper models are less efficient than shallower, wider models.The 24-layer model uses 0.3G more FLOPs than the 12-layer model.
- PlainMamba Block Design: PlainMamba’s block design achieves the best ImageNet accuracy, exceeding the matched Vision Mamba block by 3.5% and the VMamba block by 0.8%.The comparison uses identical depth, width, and training settings; the Vision Mamba block also loses performance when its CLS token is removed.
- Efficiency Comparison with ViT: For similarly sized 7.4M-parameter models, PlainMamba keeps computation low as inputs scale to 4096×4096, whereas DeiT’s FLOPs and memory rise rapidly.The DeiT-Tiny comparison model uses 224 channels, and peak memory is measured with batch size 1.
- Efficiency Comparison with ViT: At 4096×4096 resolution, DeiT token mixing reaches 23T FLOPs, 23 times the cost of PlainMamba because of self-attention’s quadratic complexity.PlainMamba’s FLOPs remain more evenly distributed across token mixing, channel mixing, and other operations.