[论文解读] U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation

Jun Ma, Li Fei-Fei+1 more|arXiv (Cornell University)|Jan 9, 2024

AI in cancer detection被引用 213

一句话总结

U-Mamba 是一种混合 CNN-SSM 网络，用于生物医学图像分割，能够线性扩展捕获长程依赖，并在不同数据集上自配置，在多项任务上超越基于 CNN 和 Transformer 的竞争对手。

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.

研究动机与目标

• 解决在生物医学图像分割中建模长程依赖性的挑战。
• 提出一个通用的 CNN-SSM 主干（U-Mamba），将局部特征提取与长程上下文结合。
• 引入自配置机制，使其能够自动适应不同数据集。
• 在四个不同的生物医学分割任务上评估 U-Mamba（3D 腹部 CT/MRI、内镜工具、显微镜细胞），并与 CNN 与 Transformer 基线进行比较。
• 提供开源代码、模型和数据，以促进重复性和进一步的研究。

提出的方法

• 引入 Mamba：选择性结构化状态空间序列模型（SSMs）结合 Scan（S6）以线性扩展捕获长程依赖，并实现输入相关的选择。
• 将 Mamba 嵌入到类似 U-Net 的编码器-解码器（U-Mamba）中，其中两个残差块位于 Mamba 块之前，用以融合局部与全局信息。
• 将 3D/2D 特征图展平成适合序列处理的形状给 Mamba 块，随后对两个平行分支进行 Hadamard 乘积融合。
• 提供两种网络变体：U-Mamba_Bot（仅瓶颈处的 Mamba）和 U-Mamba_Enc（在所有编码块中使用 Mamba）。
• 与 nnU-Net 的自配置框架整合，自动根据数据集调整网络深度和补丁大小。
• 使用 Dice 与交叉熵损失结合的训练，并采用固定的训练协议，以确保与基线的公平比较。

实验结果

研究问题

• RQ1 U-Mamba 是否能够在生物医学图像分割中超越 CNN 和 Transformer 来改善长程依赖建模？
• RQ2 线性扩展的 Mamba 块是否在多样化的二/三维生物医学分割任务中提供实际收益？
• RQ3 在腹部器官、器械和细胞分割任务中，U-Mamba 与强大的 CNN 和 Transformer 基线相比如何？
• RQ4 通过 nnU-Net 的 U-Mamba 自配置机制是否能够在无手动调参的情况下实现对不同数据集的鲁棒自适应？

主要发现

• U-Mamba_Bot 在腹部 CT 上实现 DSC 0.8683 和 NSD 0.9049，优于 CNN 与 Transformer 基线。
• U-Mamba_Enc 在 CT 上实现 DSC 0.8638 和 NSD 0.8980，在腹部 MRI 上实现 DSC 0.8501 和 NSD 0.9171，指示强的跨模态性能。
• 在 2D 任务中，U-Mamba 变体在 MRI 的器官、内镜中的器械、显微镜中的细胞等方面显示出优于基线模型的平均性能。
• 跨数据集，U-Mamba 展示出更少的分割异常和对具有挑战性的器官（如肝、胆囊）及器械的更好界定。