[论文解读] LKM-UNet: Large Kernel Vision Mamba UNet for Medical Image Segmentation

Jinhong Wang, Jintai Chen+2 more|arXiv (Cornell University)|Mar 12, 2024

Medical Image Segmentation Techniques被引用 4

一句话总结

本文提出 LMa-UNet，一种基于大感受野的 Mamba U-Net，用于 2D 和 3D 医疗图像分割，利用 Mamba 的线性复杂度，通过大窗口 SSM 实现大感受野。通过引入具有像素级和补丁级 SSM 的分层双向 Mamba 模块（LM 模块），实现了优越的局部与全局特征建模，在多器官分割基准测试中优于 CNN、Transformer 和先前的 Mamba 基础模型，DSC 和 NSD 分数均有所提升。

ABSTRACT

In clinical practice, medical image segmentation provides useful information on the contours and dimensions of target organs or tissues, facilitating improved diagnosis, analysis, and treatment. In the past few years, convolutional neural networks (CNNs) and Transformers have dominated this area, but they still suffer from either limited receptive fields or costly long-range modeling. Mamba, a State Space Sequence Model (SSM), recently emerged as a promising paradigm for long-range dependency modeling with linear complexity. In this paper, we introduce a Large Kernel Vision Mamba U-shape Network, or LKM-UNet, for medical image segmentation. A distinguishing feature of our LKM-UNet is its utilization of large Mamba kernels, excelling in locally spatial modeling compared to small kernel-based CNNs and Transformers, while maintaining superior efficiency in global modeling compared to self-attention with quadratic complexity. Additionally, we design a novel hierarchical and bidirectional Mamba block to further enhance Mamba's global and neighborhood spatial modeling capability for vision inputs. Comprehensive experiments demonstrate the feasibility and the effectiveness of using large-size Mamba kernels to achieve large receptive fields. Codes are available at https://github.com/wjh892521292/LKM-UNet.

研究动机与目标

• 为解决 CNN 和 Transformer 在医疗图像分割中的局限性，特别是其感受野受限以及长程建模计算成本高的问题。
• 探索 Mamba（一种结构化状态空间模型，SSM）在实现具有线性复杂度的大感受野空间建模方面的潜力。
• 通过设计一种分层且双向的 SSM 架构，克服 Mamba 的单向性及其在视觉任务中缺乏位置感知的问题。
• 证明基于大窗口的 Mamba 模块能够有效建模 2D 和 3D 医疗图像中的局部与全局依赖关系。
• 通过在多器官分割数据集上进行全面的消融实验与基准测试，验证所提出 LMa-UNet 的有效性。

提出的方法

• 提出一种基于大窗口的 Mamba U-Net（LMa-UNet），用大窗口 SSM 替代标准卷积和自注意力模块，以实现大感受野。
• 提出一种新颖的分层双向 Mamba 模块（LM 模块），由两种 SSM 操作组成：用于局部邻域建模的像素级 SSM（PiM）和用于全局依赖学习的补丁级 SSM（PaM）。
• 使用 PiM 缓解因分词后空间不连续性导致的 SSM 忘记问题，增强局部特征表示能力。
• 使用 PaM 以线性复杂度建模跨补丁的长程依赖关系，避免自注意力机制的二次方计算成本。
• 采用双向 Mamba（BiM）提升位置感知能力，减少序列建模中的输入顺序偏差。
• 在 nnU-Net 框架中使用 PyTorch 实现模型，各阶段采用自适应窗口大小，并使用标准训练超参数（Adam 优化器、权重衰减、初始学习率）。

实验结果

研究问题

• RQ1 与小卷积核 CNN 和窗口受限的 Transformer 相比，基于 Mamba 且采用大窗口尺寸的模型是否能在医疗图像分割中实现更优的局部与全局特征建模？
• RQ2 LM 模块中 PiM 与 PaM 的分层设计在多大程度上提升了 Mamba 在视觉任务中的表征能力？
• RQ3 双向 Mamba 在多大程度上缓解了单向 SSM 固有的位置偏差与序列顺序敏感性问题？
• RQ4 在基于 Mamba 的模型中增大窗口尺寸是否能在医疗图像分割中带来可测量的性能提升？
• RQ5 PiM、PaM 和 BiM 各组件对 LMa-UNet 整体性能的相对贡献如何？

主要发现

• 在 Abdomen CT 数据集上，LMa-UNet 的 DSC 达到 86.82，NSD 达到 90.02，优于基于 CNN、Transformer 和先前 Mamba 基础的模型。
• 在 Abdomen MR 数据集上，LMa-UNet 的 DSC 为 77.35，NSD 为 83.80，展示了在所有评估基准上的最先进性能。
• 当窗口尺寸从 [10,5,5,5,5,5,5] 增加到 [40,20,20,10,10,5,5] 时，DSC 从 75.89 提升至 77.35，NSD 从 82.26 提升至 83.80，证实了大感受野的优势。
• 消融实验表明，PiM 对性能提升的贡献大于 PaM，表明增强局部感受野比单纯加强全局建模更为关键。
• 包含全部组件（PiM + PaM + BiM）的完整 LMa-UNet 取得最佳性能，DSC 为 77.35，NSD 为 83.80，验证了分层与双向设计的有效性。
• 由于 Mamba 的线性复杂度，模型保持了高效率，实现了标准 CNN 和 Transformer 无法实现的大窗口建模。