Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

This work addresses the challenge of modeling long-range dependencies in medical image segmentation by introducing Swin-UMamba, a Mamba-based encoder–decoder framework that leverages ImageNet pretraining. The model features a VSS block with SS2D for four-directional scanning, an ImageNet-pretrained five-stage encoder, and a Swin-inspired decoder, with a lighter Swin-UMamba dagger variant using a Mamba-based decoder. Across AbdomenMRI, Endoscopy, and Microscopy datasets, ImageNet pretraining yields substantial gains and the Mamba-based approach outperforms CNNs, ViTs, and prior Mamba architectures, with Swin-UMamba dagger achieving strong results at lower computational cost. The findings highlight data efficiency, stable convergence, and reduced compute when employing ImageNet-based pretraining for Mamba-based vision models in 2D medical image segmentation, suggesting practical benefits for clinical deployment.

Abstract

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%.

Swin-UMamba: Mamba-based UNet with ImageNet-based pretraining

TL;DR

This work addresses the challenge of modeling long-range dependencies in medical image segmentation by introducing Swin-UMamba, a Mamba-based encoder–decoder framework that leverages ImageNet pretraining. The model features a VSS block with SS2D for four-directional scanning, an ImageNet-pretrained five-stage encoder, and a Swin-inspired decoder, with a lighter Swin-UMamba dagger variant using a Mamba-based decoder. Across AbdomenMRI, Endoscopy, and Microscopy datasets, ImageNet pretraining yields substantial gains and the Mamba-based approach outperforms CNNs, ViTs, and prior Mamba architectures, with Swin-UMamba dagger achieving strong results at lower computational cost. The findings highlight data efficiency, stable convergence, and reduced compute when employing ImageNet-based pretraining for Mamba-based vision models in 2D medical image segmentation, suggesting practical benefits for clinical deployment.

Abstract

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%.
Paper Structure (14 sections, 2 equations, 3 figures, 3 tables)

This paper contains 14 sections, 2 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Method
  3. Mamba-based VSS block
  4. Integrating ImageNet-based pretraining
  5. Swin-UMamba decoder
  6. Swin-UMamba$\dagger$: Swin-UMamba with Mamba-based decoder
  7. Experiments
  8. Datasets
  9. Implemetation details
  10. Baselines and evaluation metrics
  11. Comparisons on AbdomenMRI dataset
  12. Comparisons on Endoscopy dataset
  13. Comparisons on Microscopy dataset
  14. Conclusion

Figures (3)

  • Figure 1: The overall architecture of Swin-UMamba. Swin-UMamba can leverage the power of vision foundation models by loading the weights of pretrained models. Each block within the blue box was initialized with the ImageNet pretrained weights.
  • Figure 2: The overall architecture of Swin-UMamba$\dagger$.
  • Figure 3: Result visualization on a) AbdomenMRI, b) Endoscopy, and c) Microscopy. Swin-UMamba accurately recognizes the shape and type of the segmented targets.