Source-linked AI summary
Swin-UMamba: Mamba-based UNet with ImageNet-based pretraining
Jiarun Liu, Hao Yang, Hong-Yu Zhou, Yan Xi, Lequan Yu, Yizhou Yu, Yong Liang, Guangming Shi, Shaoting Zhang, Hairong Zheng, Shanshan Wang
TL;DR
Medical image segmentation needs efficient modeling of both local features and long-range dependencies, while existing CNN, ViT, and Mamba-based approaches have important limitations. The paper introduces Swin-UMamba and a lighter variant that integrate ImageNet pretraining with Mamba-based segmentation. Across medical image datasets, Swin-UMamba outperforms CNNs, ViTs, and recent Mamba-based models, with ImageNet pretraining improving segmentation performance and training behavior.
Problem
Medical image segmentation requires long-range dependency modeling, but CNNs have local receptive fields, ViTs have quadratic attention complexity, and Mamba-based models largely lack pretraining evaluation.
Method
Swin-UMamba is a Mamba-based 2D medical image segmentation network using a generic ImageNet-pretrained encoder and a designed decoder, with Swin-UMamba† as a lighter variant.
Results
Swin-UMamba outperforms CNNs, ViTs, and recent Mamba-based models, while ImageNet pretraining improves segmentation accuracy and training stability across evaluated datasets.
Takeaways & Limitations
ImageNet-based pretraining can provide superior accuracy, stable convergence, data efficiency, reduced overfitting, and lower computational resource consumption for Mamba-based medical image segmentation.
Abstract
from arXiv · showhide
Accurate medical image segmentation demands the integration of multi-scale information, spanning from local features to global dependencies. However, it is challenging for existing methods to model long-range global information, where convolutional neural networks (CNNs) are constrained by their local receptive fields, and vision transformers (ViTs) suffer from high quadratic complexity of their attention mechanism. Recently, Mamba-based models have gained great attention for their impressive ability in long sequence modeling. Several studies have demonstrated that these models can outperform popular vision models in various tasks, offering higher accuracy, lower memory consumption, and less computational burden. However, existing Mamba-based models are mostly trained from scratch and do not explore the power of pretraining, which has been proven to be quite effective for data-efficient medical image analysis. This paper introduces a novel Mamba-based model, Swin-UMamba, designed specifically for medical image segmentation tasks, leveraging the advantages of ImageNet-based pretraining. Our experimental results reveal the vital role of ImageNet-based training in enhancing the performance of Mamba-based models. Swin-UMamba demonstrates superior performance with a large margin compared to CNNs, ViTs, and latest Mamba-based models. Notably, on AbdomenMRI, Encoscopy, and Microscopy datasets, Swin-UMamba outperforms its closest counterpart U-Mamba_Enc by an average score of 2.72%.
1 Introduction
Swin-UMamba addresses the challenge of modeling global dependencies in medical image segmentation by combining Mamba-based modeling with ImageNet-pretrained vision features. The paper reports that this design improves performance over CNN, ViT, and earlier Mamba-based approaches.
- Accurate medical image segmentation must integrate local features with global dependencies, but efficiently capturing long-range information remains challenging.
- CNNs are limited by local receptive fields, while ViTs face high computational complexity in attention-based global modeling.
- Mamba-based models offer efficient long-sequence modeling with linear or near-linear scaling while retaining long-range dependency capabilities.
- Existing Mamba-based medical segmentation models are mostly trained from scratch, leaving the impact of ImageNet pretraining unclear despite limited medical datasets.
- Swin-UMamba uses a generic pretrained-model encoder and a task-designed decoder for 2D medical image segmentation.
- ImageNet-based pretraining was experimentally verified as important for Mamba-based medical image segmentation networks.
- Swin-UMamba and its lighter variant use fewer parameters or lower FLOPs while maintaining competitive performance, and the model outperforms prior CNN, ViT, and Mamba-based methods.
2 Method
Swin-UMamba combines an ImageNet-pretrained Mamba-based encoder with U-shaped decoding, skip connections, and multi-scale segmentation outputs. Its Swin-UMamba† variant replaces the CNN-based decoder with Mamba-based components to reduce computational cost.
- Swin-UMamba architecture: Swin-UMamba uses an ImageNet-pretrained Mamba-based encoder, an up-sampling decoder, and skip connections for multi-scale segmentation.The encoder extracts features at different scales, while skip connections bridge low-level details and high-level semantics.
- Mamba-based VSS block: The VSS block applies 2D-selective-scan by unfolding image patches along four directions, processing sequences with an SSM, and merging outputs into a 2D feature map.The S6 selective-scan operator lets each scanned element interact with previously scanned samples through a compressed hidden state.
- Integrating ImageNet-based pretraining: Swin-UMamba adapts VMamba-Tiny’s pretrained multi-scale structure through five encoder stages with gradual 2× down-sampling to retain low-level details.VSS blocks and patch-merging layers use ImageNet-pretrained VMamba-Tiny weights, while the patch-embedding block is not initialized from pretraining because its patch size and input channels differ.
- Swin-UMamba decoder: The decoder enhances a U-shaped design with residual convolution blocks, dense skip connections, and segmentation heads at multiple scales.The final segmentation output is obtained with a 1 × 1 convolution after feature recovery through the decoder.
- Swin-UMamba† decoder: Swin-UMamba† replaces the CNN-based up-sample blocks with patch expanding and VSS blocks, reducing parameters from 60M to 28M and FLOPs from 68.0G to 18.9G on AbdomenMRI.The variant also removes unsupported input-image skip connections and applies deep supervision at multiple resolutions.
3 Experiments
Experiments evaluate Swin-UMamba and its Mamba-decoder variant across organ, instrument, and cell segmentation datasets using accuracy and computational metrics. The models outperform baselines, while ImageNet pretraining substantially improves results, especially on smaller datasets.
- Datasets: Swin-UMamba and Swin-UMamba† were evaluated on AbdomenMRI, Endoscopy, and Microscopy, covering organ, instrument, and cell segmentation.The datasets span different resolutions and image modalities.
- Evaluation metrics: DSC, NSD, and F1 measured segmentation performance, while parameter counts and FLOPs assessed model scale and computational burden.DSC and NSD were used for AbdomenMRI and Endoscopy; F1 was used for Microscopy instance segmentation.
- AbdomenMRI: On AbdomenMRI, both Swin-UMamba variants outperformed CNN-, transformer-, and Mamba-based baselines, with Swin-UMamba improving DSC over U-Mamba_Enc by 1.34%.All Mamba-based networks exceeded CNN and transformer baselines by at least 1% on both DSC and NSD.
- AbdomenMRI: ImageNet pretraining improved Swin-UMamba by 7.06% in DSC and 7.74% in NSD on AbdomenMRI.Pretraining also enabled faster and more stable training, requiring one-tenth as many iterations as baseline methods.
- Endoscopy: On Endoscopy, Swin-UMamba† exceeded U-Mamba_Bot by 2.43% in DSC and 2.41% in NSD, while pretraining improved Swin-UMamba by 12.84% in DSC and 12.90% in NSD.The authors suggest the smaller dataset may make pretrained models useful for mitigating overfitting.
- Microscopy: On Microscopy, the two Swin-UMamba variants surpassed all baselines by 1.99% to 20.15%, with pretraining gains of 12.45% and 7.96%, respectively.The dataset has higher resolution, fewer samples, and greater visual variation, increasing demands on long-range modeling and data efficiency.
4 Conclusion
The study shows that ImageNet-based pretraining benefits Mamba-based models for 2D medical image segmentation. Swin-UMamba and Swin-UMamba† leverage pretrained models to improve segmentation performance and efficiency.
- ImageNet-based pretraining improves Mamba-based medical image segmentation models across accuracy, convergence stability, overfitting mitigation, and data efficiency.The study also reports lower computational resource consumption.
- Swin-UMamba and Swin-UMamba† are Mamba-based segmentation models designed to leverage pretrained models.
- The findings highlight pretraining as important for improving the performance and efficiency of Mamba-based models in vision tasks.