[Paper Review] Mamba-UNet: UNet-Like Pure Visual Mamba for Medical Image Segmentation

Ziyang Wang, Jian-Qing Zheng+3 more|arXiv (Cornell University)|Feb 7, 2024

Medical Image Segmentation Techniques79 citations

TL;DR

Mamba-UNet uses pure Visual Mamba blocks within a UNet-like encoder-decoder to improve long-range feature modeling, achieving higher segmentation accuracy than UNet and Swin-UNet on MRI cardiac data.

ABSTRACT

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.

Motivation & Objective

• Motivate improved long-range dependency modeling for medical image segmentation.
• Propose a UNet-inspired architecture using Visual Mamba blocks (VSS) for encoder, bottleneck, and decoder.
• Preserve spatial details via skip connections and patch merging/expanding within a VMamba-based framework.
• Evaluate segmentation performance on a public MRI cardiac dataset and compare against baselines.

Proposed method

• Adopt a UNet-like encoder–decoder with pure Visual Mamba blocks as core building blocks.
• Represent input images as patch tokens then process through hierarchical VSS blocks with patch merging/expanding.
• Use skip connections between encoder and decoder to fuse multi-scale features.
• Load pretrained VMamba-Tiny in the encoder for improved initialization.
• Train with SGD under fixed hyperparameters and evaluate with standard segmentation metrics.

Figure 1: A brief introduction of the evolution of recent developments of UNet with incorporation of Transformer and State Space Models (SSM) for medical image segmentation.

Experimental results

Research questions

• RQ1 Can VMamba-based blocks improve long-range dependency modeling in medical image segmentation compared to conventional UNet and ViT-based approaches?
• RQ2 Does a pure VMamba UNet achieve superior segmentation accuracy while maintaining computational efficiency on MRI data?
• RQ3 What is the comparative performance of Mamba-UNet versus UNet and Swin-UNet under identical training settings?

Key findings

• Mamba-UNet achieves Dice 0.9281, IoU 0.8698, and accuracy 0.9972 on the MRI cardiac test set.
• Mamba-UNet outperforms UNet in Dice and IoU under the same hyperparameters.
• Mamba-UNet shows competitive Hausdorff Distance (HD 95%) and ASD scores compared with baselines, with HD 2.4645 and ASD 0.7677.
• Swin-UNet yields Dice 0.9188 and IoU 0.8545, both below Mamba-UNet, while UNet yields Dice 0.9248.
• Mamba-UNet demonstrates superior boundary precision as reflected in the reported metrics.

Figure 2: The architecture of Mamba-UNet, which is composed of encoder, bottleneck, decoder and skip connections. The encoder, bottleneck and decoder are all constructed based on Visual Mamba block.