LightM-UNet: Mamba Assists in Lightweight UNet for Medical Image Segmentation

📄 arXiv: 2403.05246v2 📥 PDF

作者: Weibin Liao, Yinghao Zhu, Xinyuan Wang, Chengwei Pan, Yasha Wang, Liantao Ma

分类: eess.IV, cs.CV

发布日期: 2024-03-08 (更新: 2024-03-11)

🔗 代码/项目: GITHUB

💡 一句话要点

提出LightM-UNet以解决医疗图像分割中的计算资源限制问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 医疗图像分割 轻量级模型 状态空间模型 深度学习 UNet Mamba 计算资源优化

📋 核心要点

  1. 现有的UNet及其变体在医疗图像分割中存在参数量大和计算负担重的问题，尤其是基于Transformer的模型。
  2. 本文提出LightM-UNet，通过将Mamba作为轻量级替代方案集成到UNet中，旨在解决计算资源限制带来的挑战。
  3. 实验结果显示，LightM-UNet在多个真实数据集上超越了现有方法，尤其在参数和计算成本上显著降低。

📝 摘要（中文）

UNet及其变体在医疗图像分割中被广泛应用，但基于Transformer的模型因参数众多和计算负担重而不适合移动健康应用。近期，状态空间模型（SSMs）如Mamba成为CNN和Transformer架构的竞争替代方案。本文提出Lightweight Mamba UNet（LightM-UNet），将Mamba作为轻量级替代方案集成到UNet中，旨在解决实际医疗环境中的计算资源限制问题。实验结果表明，LightM-UNet在两个真实的2D/3D数据集上超越了现有的最先进文献，尤其在与nnU-Net的比较中，LightM-UNet在分割性能上表现优越，同时参数和计算成本分别降低了116倍和21倍，展示了Mamba在模型轻量化中的潜力。

🔬 方法详解

问题定义：本文旨在解决医疗图像分割中现有UNet及其变体在计算资源和参数量方面的不足，尤其是基于Transformer的模型不适合移动健康应用的问题。

核心思路：通过引入Mamba作为轻量级替代方案，LightM-UNet在UNet框架中实现了高效的深度语义特征提取和长距离空间依赖建模，旨在降低计算复杂度。

技术框架：LightM-UNet的整体架构包括Residual Vision Mamba Layer，采用纯Mamba方式进行特征提取，具有线性计算复杂度，适用于2D和3D医疗图像数据。

关键创新：LightM-UNet的主要创新在于将Mamba集成到UNet中，显著降低了模型的参数和计算成本，与传统的CNN和Transformer架构相比，提供了更高效的解决方案。

关键设计：在设计中，LightM-UNet采用了特定的损失函数和网络结构，确保在保持高分割性能的同时，最大限度地减少了计算资源的消耗。具体参数设置和网络结构细节在实验部分进行了详细说明。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，LightM-UNet在多个真实数据集上超越了现有的最先进方法，特别是在与nnU-Net的比较中，分割性能显著提升，同时参数和计算成本分别降低了116倍和21倍，展示了Mamba在模型轻量化中的巨大潜力。

🎯 应用场景

该研究的潜在应用领域包括医疗图像分析、移动健康应用和远程医疗等。通过降低模型的计算负担，LightM-UNet能够在资源受限的设备上实现高效的图像分割，具有重要的实际价值和广泛的应用前景。

📄 摘要（原文）

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.