Efficient Mamba: Overcoming the visual limitations of Mamba with innovative structures

Abstract

Mamba models have emerged as strong competitors to Transformers due to their efficient long-sequence processing and high memory efficiency. However, their state space models (SSMs) suffer from limitations in capturing long-range dependencies, lack of channel interactions, and weak generalization in vision tasks.

To address these issues, we propose Efficient Mamba (EMB), an innovative framework that enhances SSMs while integrating convolutional neural networks (CNNs) and Transformers to mitigate their inherent drawbacks. The key contributions of EMB are as follows: (1) We introduce the TransSSM module, which incorporates feature flipping and channel shuffle to enhance channel interactions and improve generalization. Additionally, we propose the Window Spatial Attention (WSA) module for precise local feature modeling and Dual Pooling Attention (DPA) to improve global feature modeling and model stability. (2) We design the MFB-SCFB composite structure, which integrates TransSSM, WSA, Inverted Residual Block(IRBs), and convolutional attention modules to facilitate effective global–local feature interaction.

EMB achieves state-of-the-art (SOTA) performance across multiple vision tasks. For instance, on ImageNet classification, EMB-S/T/N achieves Top-1 accuracies of 78.9%, 76.3%, and 73.5%, with model sizes and FLOPs of 5.9M/1.5G, 2.5M/0.6G, and 1.4M/0.3G, respectively, when trained on a single NVIDIA 4090 GPU.

Experimental results demonstrate that EMB provides a novel paradigm for efficient vision model design, offering valuable insights for future SSM research.

Code: https://github.com/Xuwei86/EMB/tree/main.