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

TL;DR

U-Mamba is a novel biomedical image segmentation network that combines CNNs with State Space Sequence Models to effectively capture long-range dependencies, outperforming existing methods across diverse datasets.

Contribution

Introduces U-Mamba, a hybrid CNN-SSM network with self-configuring capabilities for improved long-range dependency modeling in biomedical segmentation.

Findings

Outperforms state-of-the-art CNN and Transformer models on multiple tasks.

Effective in 3D CT/MR, endoscopy, and microscopy image segmentation.

Automatically adapts to various datasets without manual tuning.

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.