[Paper Review] MambaOut: Do We Really Need Mamba for Vision?

Weihao Yu, Xinchao Wang|arXiv (Cornell University)|May 13, 2024

African history and culture studies34 citations

TL;DR

MambaOut removes the state space model (SSM) from Mamba blocks and stacks gated CNN blocks to show that SSM is unnecessary for ImageNet image classification, while suggesting SSM may help long-sequence vision tasks like detection and segmentation.

ABSTRACT

Mamba, an architecture with RNN-like token mixer of state space model (SSM), was recently introduced to address the quadratic complexity of the attention mechanism and subsequently applied to vision tasks. Nevertheless, the performance of Mamba for vision is often underwhelming when compared with convolutional and attention-based models. In this paper, we delve into the essence of Mamba, and conceptually conclude that Mamba is ideally suited for tasks with long-sequence and autoregressive characteristics. For vision tasks, as image classification does not align with either characteristic, we hypothesize that Mamba is not necessary for this task; Detection and segmentation tasks are also not autoregressive, yet they adhere to the long-sequence characteristic, so we believe it is still worthwhile to explore Mamba's potential for these tasks. To empirically verify our hypotheses, we construct a series of models named MambaOut through stacking Mamba blocks while removing their core token mixer, SSM. Experimental results strongly support our hypotheses. Specifically, our MambaOut model surpasses all visual Mamba models on ImageNet image classification, indicating that Mamba is indeed unnecessary for this task. As for detection and segmentation, MambaOut cannot match the performance of state-of-the-art visual Mamba models, demonstrating the potential of Mamba for long-sequence visual tasks. The code is available at https://github.com/yuweihao/MambaOut

Motivation & Objective

• Assess whether Mamba's state space model (SSM) is necessary for visual recognition tasks.
• Evaluate the performance of SSM-free MambaOut against visual Mamba models on ImageNet classification.
• Examine the potential benefits of SSM for long-sequence visual tasks such as object detection and semantic segmentation.

Proposed method

• Construct MambaOut by stacking Gated CNN blocks (without SSM) in a ResNet-like 4-stage hierarchy.
• Replace Mamba's SSM-based token mixer with a simple depthwise Conv-based token mixer in Gated CNN blocks.
• Train on ImageNet with DeiT-style augmentations and AdamW optimization to compare against visual Mamba models.
• Evaluate on COCO for detection/segmentation with Mask R-CNN as backbone.
• Evaluate on ADE20K for semantic segmentation using UperNet backbone.

Experimental results

Research questions

• RQ1 Is SSM necessary for ImageNet image classification when using a Mamba-like architecture?
• RQ2 Can a simpler Gated CNN/block without SSM outperform visual Mamba models on ImageNet classification?
• RQ3 Does removing SSM degrade performance for long-sequence vision tasks such as object detection and semantic segmentation?
• RQ4 Is there evidence that Mamba's advantages are limited to long-sequence or autoregressive tasks in vision?

Key findings

• MambaOut, without SSM, consistently outperforms visual Mamba models on ImageNet at multiple sizes.
• MambaOut achieves higher top-1 accuracy than LocalVMamba-S and other visual Mamba variants at similar MACs.
• On COCO and ADE20K, MambaOut does not reach the performance of state-of-the-art visual Mamba models and generically trails the best convolution-attention hybrids, indicating SSM may still aid long-sequence visual tasks.
• Overall, MambaOut supports the hypothesis that SSM is unnecessary for image classification, but suggests potential benefits of SSM for detection and segmentation tasks.