PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jun Ma, Feifei Li, Bo Wang

arXiv:2401.04722v1eess.IVcs.CVcs.LG

TL;DR

Biomedical segmentation needs efficient long-range dependency modeling because CNNs are local and Transformers are computationally expensive. U-Mamba combines convolutional feature extraction with Mamba-based state-space modeling and self-configuration, and it outperforms CNN- and Transformer-based networks across diverse segmentation tasks. The authors position it as a general-purpose backbone candidate for biomedical image segmentation.

  • Problem

    CNN locality and Transformer computational complexity limit efficient long-range dependency modeling in biomedical image segmentation.

  • Method

    U-Mamba combines convolutional residual processing with Mamba-based state-space modeling in an encoder-decoder network and automatically configures its architecture across datasets.

  • Results

    U-Mamba outperformed existing CNN- and Transformer-based segmentation networks across diverse modalities and segmentation targets, including four evaluated tasks.

  • Takeaways & Limitations

    U-Mamba is presented as a versatile candidate backbone for next-generation biomedical image segmentation networks.

  • Takeaways & Limitations

    Transformer comparisons may be affected by training-resource differences because the networks were trained with one GPU despite possible multi-GPU requirements.

Abstract

from arXiv · show

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.

1 Introduction

Biomedical image segmentation supports diagnosis, treatment planning, and disease monitoring, but manual annotation is impractical and existing CNNs and Transformers face complementary limitations in modeling long-range dependencies. U-Mamba addresses this challenge with a hybrid CNN-SSM network for 2D and 3D biomedical segmentation.

  • Biomedical segmentation supports diagnosis, cancer-microenvironment quantification, treatment planning, and disease-progression tracking.
  • Manual segmentation is time-consuming, requires extensive domain knowledge, and is impractical for routine use.
  • CNNs efficiently extract hierarchical and local features, whereas Transformers improve long-range modeling but incur quadratic self-attention costs.
  • U-Mamba combines CNN-based localized feature extraction with SSM-based long-range dependency modeling for 2D and 3D biomedical image segmentation.
  • U-Mamba uses linear feature-size scaling, self-configures across datasets, and outperforms Transformer-based networks across four diverse datasets.

2 Method

U-Mamba combines convolutional residual processing with Mamba-based state-space modeling inside an encoder-decoder architecture. Its design targets efficient long-range context capture while retaining local features, skip connections, and dataset-specific self-configuration.

  • 2 Method: U-Mamba uses an encoder-decoder structure that captures local features and long-range contexts efficiently.
  • 2 Method: The conceptual architecture places two Residual blocks before an SSM-based Mamba block in each U-Mamba building block.
  • 2.1 Mamba: Selective Structured State Space Sequence Models with a Scan (S6): SSMs map one-dimensional input sequences to outputs through a latent state governed by linear differential equations.
  • 2.1 Mamba: Selective Structured State Space Sequence Models with a Scan (S6): SSMs provide linear computational complexity per time step and parallelized training, but naïve versions can require more memory and encounter vanishing gradients.
  • 2.1 Mamba: Selective Structured State Space Models with a Scan (S6): S4 structures and initializes the state matrix to support deep sequence models with efficient long-range reasoning.
  • 2.1 Mamba: Selective Structured State Space Models with a Scan (S6): Mamba adds input-dependent parameter selection and a hardware-aware recurrent scan that scales linearly with sequence length.
  • 2.1 Mamba: Selective Structured State Space Models with a Scan (S6): Mamba’s linear scaling motivates its use for image data, where Transformer self-attention creates high computational burden for large images.
  • 2.2 U-Mamba: Marry Mamba with U-Net: U-Net convolutions capture multi-scale features but have limited long-range modeling because each kernel operates within a local receptive field.

3 Experiments and results

U-Mamba was evaluated across four publicly available biomedical imaging datasets spanning 3D and 2D organ, instrument, and cell segmentation. Under controlled nnU-Net-based training and benchmarking, it outperformed the compared CNN- and Transformer-based networks, with strong quantitative and qualitative results.

  • 3.1 Datasets: Four datasets covered abdominal organ segmentation in CT and MRI, instrument segmentation in endoscopy, and cell segmentation in microscopy.The microscopy benchmark was converted from instance segmentation to semantic segmentation for architecture benchmarking.
  • 3.2 Implementation and training protocols: U-Mamba was implemented within nnU-Net to control preprocessing and augmentation, enabling comparisons in which network architecture was the differing factor.All networks were trained from scratch for 1000 epochs on one NVIDIA A100 GPU using the same batch size.
  • 3.3 Benchmarking: Evaluation used DSC and NSD for three semantic segmentation tasks, while F1 score measured cell segmentation quality.The metric choices followed recommendations from Metrics Reloaded.
  • 3.3 Benchmarking: 0.8683 and 0.8501 were U-Mamba_Bot’s abdomen CT and U-Mamba_Enc’s abdomen MR average DSC scores, respectively, surpassing CNN- and Transformer-based networks.Visual results also showed fewer outliers, more accurate liver and stomach masks in CT, and better gallbladder boundary delineation in MRI.
  • 3.4 Quantitative and qualitative segmentation results: 0.7625, 0.6504, and 0.5607 were U-Mamba’s best average DSC for organs, best average DSC for instruments, and F1 score for cells, respectively, across the 2D tasks.Qualitative comparisons reported fewer segmentation outliers and better handling of heterogeneous appearances than the compared methods.

4 Discussion and conclusion

U-Mamba addresses long-range dependency challenges in biomedical segmentation and reports stronger performance than CNN- and Transformer-based networks across diverse settings. The discussion also identifies comparison constraints and several directions for future expansion.

  • U-Mamba outperforms existing CNN- and Transformer-based segmentation networks across diverse modalities and segmentation targets.It also produces fewer segmentation outliers for objects with heterogeneous appearances.
  • U-Mamba’s architecture simultaneously extracts multi-scale local features and captures long-range dependencies, which the authors associate with its performance gains.
  • The comparison trains all networks from scratch with identical preprocessing and data augmentation on one GPU, so Transformer performance may not reflect multi-GPU or pretraining settings.
  • U-Mamba achieves overall better scores than nnU-Net despite nnU-Net’s competitive results on most tasks.The authors note complementary organ-level performance, including for the aorta, kidneys, and inferior vena cava, suggesting possible ensemble integration.
  • Future work includes large-scale training, pretrained weights, stronger augmentation, specialized losses, region-based training, and applying U-Mamba blocks to classification and detection.
  • U-Mamba is presented as a self-configuring, flexible architecture that combines CNN local pattern recognition with Mamba global context understanding for biomedical segmentation.

Appendix

The appendix provides organ-wise segmentation result tables for 3D CT, 3D MRI, and 2D MRI models. Tables identify the best and second-best scores for each organ and metric using color highlights.

  • Table 5 reports organ-wise segmentation results for 3D models on the abdomen CT dataset.Red and blue indicate the best and second-best scores for each organ and metric, respectively.
  • Table 6 reports organ-wise segmentation results for 3D models on the abdomen MRI dataset.The table uses red and blue to mark the best and second-best scores for each organ and metric.
  • Table 7 reports organ-wise segmentation results for 2D models on the abdomen MRI dataset.The table uses red and blue to mark the best and second-best scores for each organ and metric.
Loading 2401.04722v1…