PaperlayerView exact revision on arXiv

Source-linked AI summary

MedMamba: Vision Mamba for Medical Image Classification

Yubiao Yue, Zhenzhang Li

arXiv:2403.03849v5eess.IVcs.CVcs.LG

TL;DR

Medical image classification must handle both local detail and long-range dependencies across diverse modalities, while existing CNN and ViT approaches have complementary limitations. MedMamba introduces a hybrid convolution–SSM architecture and is evaluated broadly, achieving competitive performance across most tasks. The work positions Vision Mamba as a baseline for generalized medical image classification while identifying further evaluation needs.

  • Problem

    CNNs have difficulty modeling long-range dependencies, whereas ViTs incur quadratic self-attention complexity and can be difficult to deploy with limited computational resources.

  • Method

    MedMamba combines convolutional layers for local features with state-space modeling for long-range dependencies through the SS-Conv-SSM block and grouped convolutions.

  • Results

    MedMamba achieves competitive performance across generalized medical image classification tasks evaluated on 16 datasets containing ten imaging modalities and 411,007 images.

  • Takeaways & Limitations

    MedMamba provides a generalized Vision Mamba baseline for medical image classification and supports further development of Mamba-based medical AI methods.

  • Takeaways & Limitations

    Further evaluation is needed on other imaging modalities, high-resolution medical images, explainability, inference speed, and model parameter size.

Abstract

from arXiv · show

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.

1 Introduction

Medical image classification is important for computer-aided diagnosis, but growing multimodal data and difficult visual distinctions challenge manual interpretation and existing deep learning architectures. MedMamba addresses these challenges by combining convolutional and state-space modeling for generalized medical image classification, evaluated across diverse datasets.

  • Medical image classification supports clinical diagnosis, disease treatment, and disease monitoring across modalities including CT, ultrasound, X-ray, microscopy, endoscopy, and MRI.
  • CNNs capture local features but struggle with global context and long-range dependencies, while ViTs capture long-range dependencies but can degrade local feature details.
  • State-space models offer efficient long-range dependency modeling with linear or near-linear scaling in sequence length.
  • MedMamba introduces the SS-Conv-SSM hybrid block, combining channel-split convolutional and state-space branches with channel shuffling for efficient feature extraction.
  • MedMamba is evaluated on 16 datasets spanning ten imaging modalities and 411,007 images, with comparisons against numerous state-of-the-art generic visual backbones.

2 Related work

Related work spans convolutional neural networks, vision transformers, and structured state-space models for visual representation learning. These approaches differ in their handling of local structure, long-range dependencies, computational efficiency, and adaptation to medical imaging.

  • Convolution Neural Networks: CNN research progressed from early networks such as Time Delay Neural Network and LeNet-5 to influential architectures including AlexNet, VGG, and ResNet.
  • Vision Transformers: Vision transformers divide images into patches and process them as sequences using self-attention, enabling long-range dependency modeling without a two-dimensional inductive bias.
  • Visual State Space Models: Structured state-space models address earlier computational and memory challenges through efficient representations and convolution-kernel computation.
  • MedMamba’s architecture figure uses abbreviations for batch normalization, layer normalization, linear, point-wise convolution, and depth-wise convolution.
  • Visual State Space Models: Visual Mamba research includes direction-aware and bidirectional SSM designs, while medical imaging work has explored integrating Mamba layers into encoder architectures.

3 Methods

MedMamba combines state-space modeling and convolution within a hierarchical vision architecture for medical images. Its SS2D mechanism scans feature maps in multiple directions, while SS-Conv-SSM blocks integrate global and local feature extraction.

  • State-space preliminaries: SSMs map input sequences through implicit states to outputs and can be computed with linear or near-linear scaling in sequence length.The continuous formulation uses state matrix A and projection parameters B and C, then discretizes the system for deep learning.
  • Overall architecture: MedMamba comprises patch embedding, stacked SS-Conv-SSM blocks, patch merging for down-sampling, and a feature classifier.The architecture preserves 2D structure after 4 × 4 patch embedding and builds hierarchical representations across four stages.
  • Model variants: MedMamba provides Tiny, Small, and Base variants whose parameter sizes and FLOPs are specified for a 224 × 224 input and 1000 classes.The variants follow the multi-scale configuration used by typical vision transformers.
  • 2D-selective-scan: SS2D uses four-way cross-scan expansion, selective S6 scanning, and scan merging to capture long-range dependencies with linear complexity.The four directions extend one-dimensional selective scanning to two-dimensional feature maps without affecting the receptive field.
  • SS-Conv-SSM block: SS-Conv-SSM is a lightweight dual-branch block that channel-splits features between Conv-Branch and SSM-Branch for local and global information extraction.Channel concatenation restores the channel dimension, and channel-shuffle reduces information loss between grouped-convolution channels.

4 Experiments and results

The experiments evaluate MedMamba across 16 medical-image datasets spanning multiple modalities, with standardized dataset descriptions, preprocessing, training protocols, and evaluation metrics. The study also analyzes computational complexity and model size.

  • Datasets: The datasets include smartphone skin-lesion photographs, ultrasound images, fetal screening images, chest X-rays, gastrointestinal endoscopy, otoscopy, and standardized biomedical-image collections.Examples include PAD-UFES-20, Cervical-US, Fetal-Planes-DB, CPN X-ray, Kvasir, Otoscopy2024, and MedMNIST.
  • Datasets: Table 2 reports each dataset’s size, image modality, and splitting strategy.The MedMNIST splitting strategy follows prior work.
  • Experimental setup: All images were resized to 224×224×3, then each dataset was normalized and standardized before training generic vision networks and MedMamba.The resizing and normalization procedure was used for fair network comparisons.
  • Experimental setup: Non-MedMNIST models used AdamW, cross-entropy loss, 150 epochs, batch size 64, and early stopping, while MedMNIST training followed established settings for 100 epochs with batch size 128.MedMamba-X additionally used AutoAugment to examine the effect of data augmentation.
  • Evaluation: Performance was evaluated with overall accuracy, precision, sensitivity, specificity, F1-score, and AUC, alongside FLOPs and parameter size.These metrics cover predictive performance and computational complexity.

5 Results and discussions

Across diverse medical-image benchmarks, MedMamba generally combines competitive or superior accuracy with lower computational complexity. Ablations, visualization, and robustness analyses further examine the architecture’s efficiency, interpretability, and behavior under perturbations.

  • Performance comparison: 58.8% OA and 0.808 AUC are achieved by MedMamba-T on PAD-UFES-20 with the lowest FLOPs.Its OA and AUC exceed ConvNeXt-T by 4.5% and 0.048, respectively; on Cervical-US, it achieves the best OA and exceeds VMamba-T, Swin-T, and ConvNeXt-T by 1.8%, 6.8%, and 4.0%.
  • Performance comparison: 28.1M, 39.6M, and 40.5M fewer parameters are used by MedMamba-B than VMamba-B, Swin-B, and ConvNeXt-B on Fetal-Planes-DB.MedMamba-B also has the lowest FLOPs and improves OA over those models by 0.6%, 5.2%, and 5.3%, respectively.
  • Performance comparison: 84.0%, 84.3%, and 83.8% average OA are achieved by MedMamba-T, MedMamba-S, and MedMamba-B on non-MedMNIST, respectively.These variants improve over VMamba, Swin, and ConvNeXt counterparts while achieving a favorable parameter–FLOPs–OA trade-off.
  • Performance comparison: 7.2% higher OA and 1.0% higher AUC than ResNet50 are achieved by MedMamba-S on PathMNIST.On OCTMNIST, MedMamba-S improves OA and AUC over MedViT-S by 14.7% and 3.6%, respectively.
  • Performance comparison: 86.2%, 87.0%, and 87.4% average OA are achieved by MedMamba-T, MedMamba-S, and MedMamba-B on MedMNIST, respectively.The variants improve over similarly scaled MedViT models, with substantially lower FLOPs and parameter sizes.
  • Ablation Study: 35.4M parameters and 5.4G FLOPs are added when the SSM and convolution branches are fused by direct element-wise addition.Grouped convolution via channel splitting alleviates this complexity, while channel shuffling further strengthens cross-group channel interaction and improves performance.

6 Conclusion

MedMamba is presented as the first Vision Mamba for generalized medical image classification, combining convolutional layers and SSM modules to model local features and long-range dependencies. Extensive evaluation across diverse datasets supports its competitiveness, while several application areas remain future work.

  • MedMamba combines classic convolutional layers and SSM modules to extract local representations and model long-range dependencies for generalized medical image classification.
  • 411,007 images across 16 datasets and ten imaging modalities were used to evaluate MedMamba comprehensively.
  • MedMamba achieves extremely competitive performance against state-of-the-art methods across the evaluated datasets.
  • Grouped convolution provides a trade-off between efficient modeling and computational resource consumption.
  • Future work includes testing additional imaging modalities, optimizing the architecture, analyzing decisions with explainable AI, and studying high-resolution medical images.
Loading 2403.03849v5…