[论文解读] LightM-UNet: Mamba Assists in Lightweight UNet for Medical Image Segmentation

Weibin Liao, Yinghao Zhu+4 more|arXiv (Cornell University)|Mar 8, 2024

COVID-19 diagnosis using AI被引用 46

一句话总结

LightM-UNet 将 Mamba 与 UNet 结合，创建一个轻量级的3D/2D 医学图像分割模型，参数 ~1M，在 LiTS 上达到最先进的结果同时大幅减少参数量和 FLOPs。

ABSTRACT

UNet and its variants have been widely used in medical image segmentation. However, these models, especially those based on Transformer architectures, pose challenges due to their large number of parameters and computational loads, making them unsuitable for mobile health applications. Recently, State Space Models (SSMs), exemplified by Mamba, have emerged as competitive alternatives to CNN and Transformer architectures. Building upon this, we employ Mamba as a lightweight substitute for CNN and Transformer within UNet, aiming at tackling challenges stemming from computational resource limitations in real medical settings. To this end, we introduce the Lightweight Mamba UNet (LightM-UNet) that integrates Mamba and UNet in a lightweight framework. Specifically, LightM-UNet leverages the Residual Vision Mamba Layer in a pure Mamba fashion to extract deep semantic features and model long-range spatial dependencies, with linear computational complexity. Extensive experiments conducted on two real-world 2D/3D datasets demonstrate that LightM-UNet surpasses existing state-of-the-art literature. Notably, when compared to the renowned nnU-Net, LightM-UNet achieves superior segmentation performance while drastically reducing parameter and computation costs by 116x and 21x, respectively. This highlights the potential of Mamba in facilitating model lightweighting. Our code implementation is publicly available at https://github.com/MrBlankness/LightM-UNet.

研究动机与目标

• 通过降低计算负担和参数量，推动适用于移动健康的轻量级医学图像分割。
• 利用 Mamba（状态空间模型）在 UNet 内捕捉长程依赖，而不产生 Transformer 级的成本。
• 提出残差视觉 Mamba 层（RVM 层）和视觉状态空间模块（VSS 模块），在尽量少的参数增加下增强深层特征建模。
• 在保持参数量约为1M的同时，展示在 3D LiTS 及 2D Montgomery&Shenzhen 数据集上的最先进性能。

提出的方法

• 将 Mamba 作为纯粹基于 Mamba 的编码器嵌入到类似 UNet 的架构中。
• 使用带 LayerNorm、VSSM 和残差调整因子的 Residual Vision Mamba Layer（RVM Layer）来建模长程依赖。
• 结合 Vision State-Space Module（VSS Module）具有并行分支和 Hadamard 乘积以融合局部与全局特征。
• 以 UNet 风格的编码器–瓶颈–解码器布局，将 LightM-UNet 结构为 3 个编码器块、一个瓶颈块和 3 个解码器块。
• 对浅层特征提取应用深度可卷积，并在末端使用双线性上采样以产生分割图。
• 使用 3D LiTS 和 2D Montgomery&Shenzhen 数据集，结合交叉熵和 Dice 损失进行训练，并与 nnU-Net、SegResNet、UNETR、SwinUNETR 和 U-Mamba 进行对比。

实验结果

研究问题

• RQ1 基于 UNet 的医学图像分割是否能够在轻量化架构下实现全局长程建模？
• RQ2 纯 Mamba 编码器在 3D/2D 分割的准确性和效率方面，与 CNN/Transformer 块相比如何？
• RQ3 RVM Layer 与 VSS Module 是否在几乎不增加参数的情况下提供有意义的增益？
• RQ4 LightM-UNet 在大幅减少参数与 FLOPs 的同时，是否具有与最先进模型的竞争力？

主要发现

• LightM-UNet 在 3D LiTS 任务中实现 1.87M 参数和 457.62 GFLOPs。
• 在 LiTS 上，LightM-UNet 获得肝脏 DSC 96.31 和肿瘤 DSC 72.86，平均 DSC 84.58，平均 mIoU 77.48。
• 相较于 nnU-Net，LightM-UNet 将参数量和 FLOPs 分别降低了 47.39x 和 15.82x，同时实现更高的平均 DSC/mIoU。
• 对 U-Mamba，LightM-UNet 将平均 mIoU 提高 2.11 个百分点，肿瘤 DSC 提高 3.63 个百分点。
• 消融结果显示用 CNN 或自注意力替换 VSSM 会降低性能并增加参数量，而移除调整因子或残差连接会降低性能且不节省成本。
• 总体而言，LightM-UNet 在 3D LiTS 上提供最先进的性能，并在 2D 数据集上达到具竞争力的结果，且极为轻量。