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

题目：U-Mamba: Enhancing Long-Range Dependency for Biomedical Image Segmentation

作者：Jiarun Liu, Hao Yang, Hongyu Zhou, Yan Xi, Lequan Yu, Cheng Li, Yong Xia, Yizhou Yu

链接：https://arxiv.org/pdf/2401.04722

1. Motivation (动机)

在生物医学图像分割领域，U-Net及其变体等卷积神经网络（CNN）虽然表现出色，但其固有的局部感受野限制了对全局上下文信息的捕捉。尽管Transformer能够建模长距离依赖，但其二次方的计算复杂度使其在处理高分辨率医学图像时效率低下。Mamba作为一种新兴的状态空间模型（SSM），具备线性复杂度和强大的全局建模能力，但原生Mamba是为处理1D序列数据设计的。因此，本文的核心动机是探索如何将Mamba有效地适配到2D视觉任务中，构建一个既能捕捉长距离依赖又具备高计算效率的医学图像分割模型。

2. Methods (方法)

本文提出了U-Mamba架构，其核心思想是将Mamba模块作为即插即用的组件集成到经典的U-Net框架中，以增强其全局建模能力。

核心组件：视觉Mamba块（Vision Mamba Block, VMB）
VMB是U-Mamba的基本构建单元，旨在替代U-Net中的传统卷积块。其内部结构如下：

• 线性嵌入：首先通过线性层将输入特征图的通道数进行变换。
• 深度卷积：使用深度卷积（Depthwise Convolution）来编码局部空间信息，弥补SSM在局部特征提取上的不足。
• SiLU激活：对局部特征进行非线性变换。
• Mamba层：这是模块的核心，负责捕捉长距离依赖和全局上下文信息。
• 残差连接：将输入特征与经过Mamba层处理后的特征相加，以稳定训练并促进信息流动。

网络架构：U-Net与VMB的结合
U-Mamba的整体架构遵循U-Net的编码器-解码器设计。

• 编码器：通过下采样路径逐步提取特征。在网络的较深层，作者用VMB替换了原有的卷积块，利用Mamba的全局感受野来整合上下文信息。
• 解码器：通过上采样路径逐步恢复空间分辨率，并结合跳跃连接（Skip Connections）融合来自编码器的对应层特征，以实现精确的像素级分割。
• 混合设计：这种设计巧妙地结合了CNN在浅层提取局部细节的优势和Mamba在深层建模全局依赖的能力。

3. Experiment (实验)

作者在多个权威的医学图像分割数据集上验证了U-Mamba的有效性，包括Synapse（多器官CT分割）、ACDC（心脏MRI分割）和MoNuSeg（细胞核分割）。

• 对比实验：U-Mamba与多种主流模型进行了对比，包括基于CNN的U-Net、UNet++，以及基于Transformer的TransUNet、Swin-UNet等。
• 评价指标：主要采用Dice系数（DSC）和Hausdorff距离（HD95）来评估分割精度。
• 结果分析：
  • 在Synapse和ACDC数据集上，U-Mamba均取得了优于对比方法的性能，证明了其在处理复杂解剖结构时的优势。
  • 消融实验证实，将VMB集成到U-Net中能显著提升模型性能，验证了Mamba模块在增强长距离依赖建模方面的有效性。
  • 与Transformer模型相比，U-Mamba在保持高性能的同时，展现了更优的计算效率和更低的内存占用。

4. 总结

本文提出的U-Mamba成功地将状态空间模型（Mamba）引入到2D医学图像分割任务中。通过设计视觉Mamba块（VMB）并将其融入U-Net架构，U-Mamba有效地结合了CNN的局部特征提取能力和Mamba的全局上下文建模能力，实现了线性复杂度下的高精度分割。实验结果表明，U-Mamba在多个基准测试中均超越了现有的CNN和Transformer模型，为高效、精确的医学图像分析提供了一种新的有力工具。