arXiv 2403.052462024-03-11eess.IVcs.CV

LightM-UNet：Mamba助力用于医学图像分割的轻量化UNet

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

School of Computer Science, Peking University（北京大学计算机学院）；Institute of Artificial Intelligence, Beihang University（北京航空航天大学人工智能研究院）；Key Laboratory of High Confidence Software Technologies, Ministry of Education（教育部高可信软件技术重点实验室）；National Engineering Research Center for Software Engineering, Peking University（北京大学国家软件工程研究中心）

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

阅读 PDF

arXiv 页面 ↗HTML 原文 ↗

AI总结：

针对医学图像分割中UNet类模型参数量大、计算负载高的问题，提出融合Mamba的轻量化LightM-UNet，以纯Mamba结构提取特征建模长程依赖，在2D/3D数据集上性能超前沿方法，大幅降低参数与计算成本。

摘要：

UNet及其变体已被广泛应用于医学图像分割领域。然而，这些模型，尤其是基于Transformer架构的模型，因参数量庞大、计算负载高而带来诸多挑战，不适用于移动健康应用。近期，以Mamba为代表的状态空间模型（State Space Models, SSMs）成为CNN和Transformer架构的有力替代方案。基于此，我们将Mamba用作UNet中CNN和Transformer的轻量化替代方案，旨在解决实际医疗场景中计算资源受限带来的挑战。为此，我们提出了轻量化Mamba UNet（LightM-UNet），在轻量化框架中整合了Mamba与UNet。具体而言，LightM-UNet以纯Mamba的方式利用残差视觉Mamba层提取深层语义特征并建模长程空间依赖，且具有线性计算复杂度。在两个真实世界2D/3D数据集上开展的大量实验表明，LightM-UNet的性能优于现有前沿研究成果。值得注意的是，与知名的nnU-Net相比，LightM-UNet实现了更优的分割性能，同时参数量和计算成本分别大幅降低116倍和21倍。这凸显了Mamba在推动模型轻量化方面的潜力。我们的代码实现已公开于https://github.com/MrBlankness/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.

机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。