U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation

Abstract: Convolutional Neural Networks (CNNs) and Transformers have been the most popular architectures for biomedical image segmentation, but both of them have limited ability to handle long-range dependencies because of inherent locality or computational complexity. To address this challenge, we introduce U-Mamba, a general-purpose network for biomedical image segmentation. Inspired by the State Space Sequence Models (SSMs), a new family of deep sequence models known for their strong capability in handling long sequences, we design a hybrid CNN-SSM block that integrates the local feature extraction power of convolutional layers with the abilities of SSMs for capturing the long-range dependency. Moreover, U-Mamba enjoys a self-configuring mechanism, allowing it to automatically adapt to various datasets without manual intervention. We conduct extensive experiments on four diverse tasks, including the 3D abdominal organ segmentation in CT and MR images, instrument segmentation in endoscopy images, and cell segmentation in microscopy images. The results reveal that U-Mamba outperforms state-of-the-art CNN-based and Transformer-based segmentation networks across all tasks. This opens new avenues for efficient long-range dependency modeling in biomedical image analysis. The code, models, and data are publicly available at https://wanglab.ai/u-mamba.html.

• The paper presents a novel hybrid architecture combining CNNs with state space models to capture long-range dependencies in biomedical image segmentation.
• The methodology integrates U-Mamba blocks into a self-configuring encoder-decoder framework, achieving superior DSC scores across various imaging modalities.
• Experimental results show improved accuracy and reduced segmentation outliers compared to conventional CNN and Transformer models.

U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation

The paper "U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation" presents a novel segmentation network designed to address limitations in long-range dependency modeling inherent in current CNN and Transformer architectures. It introduces an innovative hybrid architecture that combines CNNs with State Space Sequence Models (SSMs), specifically the Mamba block, to enhance segmentation in biomedical imaging tasks.

Introduction to U-Mamba

U-Mamba is grounded in the need for efficient long-range dependency modeling in biomedical image segmentation, a requirement that existing architectures struggle to meet either due to their locality (as in CNNs) or computational complexity (as in Transformers). The U-Mamba architecture effectively integrates the hierarchical feature extraction capabilities of CNNs with the superior long-sequence modeling of SSMs. This integration facilitates capturing both local and global features, offering improved segmentation performance across various biomedical imaging modalities.

Figure 1: Overview of the U-Mamba (Enc) architecture. It highlights the use of U-Mamba blocks and their integration within an encoder-decoder framework with skip connections.

Architectural Details

U-Mamba leverages a self-configuring encoder-decoder structure, wherein the encoder benefits from U-Mamba blocks composed of two Residual blocks followed by a Mamba block, as illustrated in Figure 1. The Mamba block, inspired by structured state space models (S4), selectively handles relevant information from input data, employing an efficient hardware-aware implementation that ensures linear scaling with sequence length. This is particularly beneficial given the high resolution and complexity typical in biomedical images.

The self-configuring feature, inherited from nnU-Net, ensures adaptability across different datasets, further enhancing U-Mamba's versatility. The encoder effectively captures long-range dependencies, while the decoder focuses on refining local details, all facilitated through strategic skip connections.

Experimental Results

The experimental evaluation demonstrates the superiority of U-Mamba over traditional CNN-based architectures (nnU-Net, SegResNet) and Transformer-based networks (UNETR, SwinUNETR) across multiple datasets, including 3D abdominal organ segmentation in CT and MRI scans, instrument segmentation in endoscopy images, and cell segmentation in microscopy images.

Figure 2: Visualized segmentation examples of abdominal organ segmentation in CT (1st and 2nd rows) and MRI scans (3rd and 4th rows), highlighting U-Mamba's capability to distinguish complex soft tissues.

U-Mamba consistently achieved higher Dice Similarity Coefficient (DSC) scores, outperforming competitors in robustness to heterogeneous appearances and reducing segmentation outliers, as evidenced in Figures 2 and 3.

Performance and Comparisons

Quantitative results, as summarized in Table \ref{tab:results-3d} and \ref{tab:results-2d}, underscore U-Mamba's advantage in processing both 3D and 2D data. The two network variants, U-Mamba_Bot and U-Mamba_Enc, effectively surpass existing models, achieving notable improvements in metrics such as DSC for organ segmentation tasks (Figure 3).

Figure 3: Visualized examples of MRI organ segmentation, cell segmentation, and endoscopy instrument segmentation, reinforcing U-Mamba's robustness.

Discussion and Future Work

The introduction of U-Mamba marks significant progress in biomedical image segmentation, particularly through its efficient handling of long-range dependencies. The paper suggests that the integration of CNN and SSMs could pave the way for more scalable, flexible, and robust networks. Future directions include leveraging large-scale pre-training for enhanced model transferability, optimizing loss functions for imbalanced targets, and integrating with existing classification and detection paradigms, potentially expanding U-Mamba's versatility beyond its current applications.

Conclusion

U-Mamba successfully marries convolutional operations with Mamba state space models to offer a scalable solution for biomedical image segmentation. Its self-configuring capability and performance superiority suggest a promising future for this architecture as a backbone in next-generation biomedical imaging solutions.