Source-linked AI summary
Mamba-UNet: UNet-Like Pure Visual Mamba for Medical Image Segmentation
Ziyang Wang, Jian-Qing Zheng, Yichi Zhang, Ge Cui, Lei Li
TL;DR
Medical image segmentation needs efficient modeling of long-range dependencies, while Transformer self-attention can be computationally costly for high-resolution images. The paper proposes Mamba-UNet, a pure Visual Mamba U-Net with hierarchical encoder-decoder paths and skip connections, and reports superior performance to classical UNet and Swin-UNet on the evaluated datasets.
Problem
High-resolution medical image segmentation requires long-range dependency modeling, but Transformer self-attention scales quadratically with input size.
Method
Mamba-UNet uses a pure Visual Mamba encoder-decoder with VSS blocks, hierarchical feature processing, patch merging and expanding, and skip connections.
Results
Mamba-UNet demonstrates superior performance against classical UNet and similar networks such as Swin-UNet on the evaluated medical image segmentation datasets.
Takeaways & Limitations
The architecture provides a U-Net-style segmentation network based entirely on Visual Mamba blocks while preserving multiscale spatial information through skip connections.
Abstract
from arXiv · showhide
In recent advancements in medical image analysis, Convolutional Neural Networks (CNN) and Vision Transformers (ViT) have set significant benchmarks. While the former excels in capturing local features through its convolution operations, the latter achieves remarkable global context understanding by leveraging self-attention mechanisms. However, both architectures exhibit limitations in efficiently modeling long-range dependencies within medical images, which is a critical aspect for precise segmentation. Inspired by the Mamba architecture, known for its proficiency in handling long sequences and global contextual information with enhanced computational efficiency as a State Space Model (SSM), we propose Mamba-UNet, a novel architecture that synergizes the U-Net in medical image segmentation with Mamba's capability. Mamba-UNet adopts a pure Visual Mamba (VMamba)-based encoder-decoder structure, infused with skip connections to preserve spatial information across different scales of the network. This design facilitates a comprehensive feature learning process, capturing intricate details and broader semantic contexts within medical images. We introduce a novel integration mechanism within the VMamba blocks to ensure seamless connectivity and information flow between the encoder and decoder paths, enhancing the segmentation performance. We conducted experiments on publicly available ACDC MRI Cardiac segmentation dataset, and Synapse CT Abdomen segmentation dataset. The results show that Mamba-UNet outperforms several types of UNet in medical image segmentation under the same hyper-parameter setting. The source code and baseline implementations are available.
1 Introduction
Medical image segmentation builds on U-Net’s multiscale encoder-decoder and skip connections, while newer Transformer and State Space Model approaches target long-range feature dependencies. Mamba-UNet combines a Visual Mamba-based architecture with U-Net structure for this setting.
- U-Net uses symmetric encoder-decoder paths and skip connections to extract and efficiently transform features across multiple levels.
- Transformer self-attention has quadratic computational scaling with input size, creating challenges for high-resolution biomedical images.
- State Space Models offer efficient long-sequence processing, while Mamba adds selective mechanisms and hardware optimization for dense data.
- Mamba-UNet uses a Visual Mamba-based encoder-decoder with skip connections, integrating State Space Models into a U-Net-like segmentation architecture.
2 Approach
Mamba-UNet is a hierarchical encoder-decoder that processes image patches with Visual State Space blocks, downsampling and upsampling features while using skip connections to preserve spatial detail. Its VSS blocks combine convolutional and State Space processing for feature learning.
- 2.1 Architecture Overview: The model converts a 2D grayscale image into patch tokens, embeds them, and processes them through VSS blocks and patch merging to form hierarchical features.
- 2.1 Architecture Overview: The decoder uses VSS blocks and patch expanding layers to upscale deep features by 2× while halving feature dimensions.
- 2.2 VSS Block: The underlying SSM maps input x(t) to output y(t) through hidden state h(t), with A, B, C, and D serving as evolution, projection, and skip parameters.
- 2.2 VSS Block: The discrete linear model is obtained through zero-order hold using a timescale parameter ∆, with B and C represented in the discretized formulation.
- 2.2 VSS Block: Each VSS block applies linear embedding and dual pathways, including depth-wise convolution, SiLU activation, SS2D processing, normalization, and feature merging.
- 2.1 Architecture Overview: Encoder stages maintain feature resolution through VSS blocks, while repeated patch merging reduces token count by 1/2 and doubles feature dimensions.
- 2.1 Architecture Overview: Skip connections blend encoder and decoder features at each level, combining multiscale representations with upscaled outputs to enhance spatial detail.
3 Experiments and Results
Experiments evaluate Mamba-UNet on the ACDC cardiac MRI dataset against several segmentation baselines using standardized training settings, multiple metrics, qualitative comparisons, and Dice-based prediction distributions.
- Datasets: The ACDC dataset contains MRI scans from 100 patients covering multiple cardiac structures and five pathological subgroups.The annotated structures include the right ventricle and the endocardial and epicardial walls of the left ventricle.
- Experimental Setup: Training and inference required approximately 5 hours for ACDC using 10,000 iterations and a batch size of 24.The reported runtime includes data transfer, model training, and inference.
- Experimental Setup: All baseline methods were trained with identical hyperparameter configurations for direct comparison with Mamba-UNet.The compared methods include UNet, Attention UNet, TransUNet, and Swin-UNet.
- Evaluation Metrics: Evaluation includes higher-is-better similarity metrics and lower-is-better surface-distance metrics for comparing predicted and ground-truth segmentations.Similarity measures include Dice, IoU, Accuracy, Precision, Sensitivity, and Specificity; difference measures include HD95 and ASD.
- Qualitative Results: Qualitative evaluation compares three randomly selected raw images, ground truths, and predictions from all baseline methods on the ACDC dataset.Figure 4 presents the visual comparison of UNet, Attention UNet, TransUNet, Swin-UNet, and Mamba-UNet outputs.
- Quantitative Results: Quantitative results across two datasets indicate that Mamba-UNet more likely produces precise segmentation masks and high-Dice predictions.Figure 6 shows the distribution of image-by-image Dice coefficients, with the x-axis representing Dice coefficient and the y-axis representing prediction count.
4 Conclusion
Mamba-UNet is presented as a purely Visual Mamba block-based U-Net-style network for medical image segmentation. It demonstrates superior performance against similar networks, while future work targets broader tasks, modalities, 3D images, and weaker supervision.
- The method is compared visually with other segmentation methods against ground truth on the Synapse CT Abdomen dataset.
- Direct performance comparisons are reported for segmentation networks on the ACDC MRI Cardiac test set.
- Dice distributions are visualized for Mamba-UNet and other methods on both the ACDC MRI Cardiac and Synapse CT Abdomen datasets.
- Direct performance comparisons are also reported for segmentation networks on the Synapse CT abdominal test set.
- Future work will explore more medical image segmentation tasks, different modalities and targets, additional backbones, 3D images, and semi/weakly supervised learning.