PlainMamba: Improving Non-Hierarchical Mamba in Visual Recognition

TL;DR

PlainMamba introduces a simple, scalable non-hierarchical state space model for visual recognition, leveraging continuous 2D scanning and direction-aware updates to improve spatial feature learning and efficiency.

Contribution

It adapts the Mamba state space model to images with novel 2D scanning and directional encoding, simplifying architecture and enhancing performance on visual tasks.

Findings

Achieves performance gains over previous non-hierarchical models

Competitive with hierarchical models on various visual recognition tasks

Requires less computation for high-resolution inputs

Abstract

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 .