【论文解读】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.047221. Motivation (动机)在生物医学图像分割领域U-Net及其变体等卷积神经网络CNN虽然表现出色但其固有的局部感受野限制了对全局上下文信息的捕捉。尽管Transformer能够建模长距离依赖但其二次方的计算复杂度使其在处理高分辨率医学图像时效率低下。Mamba作为一种新兴的状态空间模型SSM具备线性复杂度和强大的全局建模能力但原生Mamba是为处理1D序列数据设计的。因此本文的核心动机是探索如何将Mamba有效地适配到2D视觉任务中构建一个既能捕捉长距离依赖又具备高计算效率的医学图像分割模型。2. Methods (方法)本文提出了U-Mamba架构其核心思想是将Mamba模块作为即插即用的组件集成到经典的U-Net框架中以增强其全局建模能力。核心组件视觉Mamba块Vision Mamba Block, VMBVMB是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模型为高效、精确的医学图像分析提供了一种新的有力工具。