[Paper Review] MedMamba: Vision Mamba for Medical Image Classification

Yubiao Yue, Zhenzhang Li|arXiv (Cornell University)|Mar 6, 2024

COVID-19 diagnosis using AI65 citations

TL;DR

MedMamba introduces a Vision Mamba-based architecture with an SS-Conv-SSM block to capture both local and long-range features for medical image classification, evaluated across multiple modalities with competitive results.

ABSTRACT

Since the era of deep learning, convolutional neural networks (CNNs) and vision transformers (ViTs) have been extensively studied and widely used in medical image classification tasks. Unfortunately, CNN's limitations in modeling long-range dependencies result in poor classification performances. In contrast, ViTs are hampered by the quadratic computational complexity of their self-attention mechanism, making them difficult to deploy in real-world settings with limited computational resources. Recent studies have shown that state space models (SSMs) represented by Mamba can effectively model long-range dependencies while maintaining linear computational complexity. Inspired by it, we proposed MedMamba, the first Vision Mamba for generalized medical image classification. Concretely, we introduced a novel hybrid basic block named SS-Conv-SSM, which purely integrates the convolutional layers for extracting local features with the abilities of SSM to capture long-range dependencies, aiming to model medical images from different image modalities efficiently. By employing the grouped convolution strategy and channel-shuffle operation, MedMamba successfully provides fewer model parameters and a lower computational burden for efficient applications without sacrificing accuracy. We thoroughly evaluated MedMamba using 16 datasets containing ten imaging modalities and 411,007 images. Experimental results show that MedMamba demonstrates competitive performance on most tasks compared with the state-of-the-art methods. This work aims to explore the potential of Vision Mamba and establish a new baseline for medical image classification, thereby providing valuable insights for developing more powerful Mamba-based artificial intelligence algorithms and applications in medicine. The source codes and all pre-trained weights of MedMamba are available at https://github.com/YubiaoYue/MedMamba.

Motivation & Objective

• Address the limitation of CNNs in capturing long-range dependencies in medical images.
• Propose a lightweight, efficient architecture combining local and long-range modeling.
• Evaluate MedMamba on diverse medical imaging modalities to establish a new baseline for SSM-based medical image classification.

Proposed method

• Introduce SS-Conv-SSM module that fuses convolutional local features with state-space model long-range dependencies.
• Develop SS2D (2D-selective-scan) to scan and merge multi-directional feature sequences.
• Use a patch embedding layer (4x4 patches) followed by four stages of SS-Conv-SSM blocks with channel expansion.
• Apply layer normalization, dual-branch architecture (Conv-branch and SSM-branch) with channel shuffle for interaction.
• Train without data augmentation or pre-training to demonstrate architecture-driven gains.]
• Provide the theoretical basis via discretized state-space modeling (Abar, Bbar) and convolutional kernel assembly (Kbar) for end-to-end modeling of image data.

Experimental results

Research questions

• RQ1 How does integrating state-space models with convolutional features affect medical image classification across modalities?
• RQ2 Can MedMamba achieve competitive performance without pre-training or data augmentation on diverse medical datasets?
• RQ3 What is the impact of the SS2D selective scanning mechanism on feature extraction and long-range dependency modeling?

Key findings

• MedMamba achieves competitive and, on several datasets, superior performance compared with CNN- and Transformer-based baselines.
• The architecture performs well across multiple modalities, including skin lesions, ultrasound, chest X-ray, endoscopy, and MedMNIST-derived tasks.
• SS-Conv-SSM effectively combines local (convolution) and global (SSM) features to capture fine-grained and coarse-grained information in medical images.
• Extensive experiments on five public/private medical datasets suggest MedMamba can establish a new baseline for SSM-based medical image classification.
• Source-code is publicly available for replication and further research.