PlainMamba: Improving Non-Hierarchical Mamba in Visual Recognition
Chenhongyi Yang Zehui Chen Miguel Espinosa Linus Ericsson Zhenyu Wang Jiaming Liu Elliot J. Crowley
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 including image classification, semantic
segmentation, object detection, and instance segmentation. Our method achieves
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