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

2024-03-06

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

Image and Video Processing,Computer Vision and Pattern Recognition,Machine Learning

What problem does this paper attempt to address?

The paper attempts to address the issue of exploring the effectiveness of the Mamba model based on ImageNet pre-training in 2D medical image segmentation tasks and proposes a new Mamba model, Swin-UMamba, and its variant Swin-UMamba†.

Specifically, the paper points out the limitations of existing Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs) in handling local features and global dependencies. The Mamba model excels in long-sequence modeling, but most existing Mamba models are trained from scratch, not leveraging the advantages of pre-training. Therefore, the authors propose a new Mamba model, Swin-UMamba, which combines the benefits of ImageNet pre-training and designs an encoder-decoder structure specifically for medical image segmentation tasks.

The main contributions of the paper include:
1. For the first time, exploring the impact of pre-trained Mamba models in medical image segmentation, verifying the importance of ImageNet pre-training for Mamba models.
2. Proposing Swin-UMamba and its variant Swin-UMamba†, which has fewer parameters and lower computational complexity.
3. Experimental results show that Swin-UMamba and its variant significantly outperform existing CNNs, ViTs, and the latest Mamba models on multiple datasets, especially on small datasets.

Through these studies, the authors demonstrate the important role of ImageNet pre-training in improving the performance of medical image segmentation tasks and propose an efficient and accurate new model.