LightM-UNet：Mamba轻量级UNet进行医学图像分割

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

单位：北京大学, 北航

代码：https://github.com/MrBlankness/LightM-UNet

论文：https://arxiv.org/abs/2403.05246

CVPR 2024 论文和开源项目合集请戳—>https://github.com/amusi/CVPR2024-Papers-with-Code

目录

1、摘要

2、引言

3、模型

4、实验：

1、摘要

UNet及其变体已成功地应用于医学图像分割中，并取得了较好的效果.然而，这些模型，尤其是基于Transformer架构的模型，由于其大量参数和计算负载而带来了挑战，使得它们不适合移动的健康应用。最近，以Mamba为代表的状态空间模型（SSM）已经成为CNN和Transformer架构的竞争性替代方案。在此基础上，我们在UNet中使用了Mamba作为CNN和Transformer的轻量级替代品，旨在解决真实的医疗环境中计算资源限制所带来的挑战。为此，我们引入了轻量级的Mamba UNet（LightM-UNet），它将Mamba和UNet集成在一个轻量级的框架中。具体而言，LightM-UNet以纯Mamba的方式利用残差视觉Mamba层，以线性计算复杂度提取深层语义特征并对长范围空间依赖性进行建模。在两个真实世界的2D/3D数据集上进行的大量实验表明，LightM-UNet超越了现有的最先进的文献。值得注意的是，与著名的nnU-Net相比，LightM-UNet在实现上级的分割性能的同时，将参数和计算成本分别大幅降低了116倍和21倍。这突出了Mamba在促进模型轻量化方面的潜力。

Mamba,UNet,轻量级，卷积的架构

2、引言

UNet [16]作为一种成熟的医学图像分割算法，广泛应用于与医学器官和病变相关的一系列分割任务，涵盖各种医学图像模式。其对称的U形编码器-解码器架构，加上完整的跳过连接，为分割模型奠定了基础，产生了大量的U型架构模型。.然而，作为一种基于卷积神经网络（CNN）的模型，UNet努力解决卷积运算的固有局部性，这限制了其理解显式全局和远程语义信息交互的能力[2]。一些研究试图通过采用atrous卷积层[5]，自我注意机制[19]和图像金字塔[25]来缓解这个问题。尽管如此，这些方法仍然表现出在建模长期依赖性的约束。

为了赋予UNet理解全球信息的能力，最近的研究[2，7，6]已经深入研究了集成Transformer架构[22]，利用自我注意力机制通过将图像视为连续的补丁序列来捕获全球信息。虽然有效，但基于transformer-based的解决方案由于自注意机制而引入了关于图像大小的二次复杂性，导致大量的计算开销，特别是对于需要密集预测的任务，例如医学图像分割。这忽略了现实世界医疗环境中计算约束的必要性，无法满足移动的医疗保健分割任务中以低参数和最小计算负载为特征的模型的必要性[18]。总之，未解决的问题仍然存在：“如何才能赋予UNet适应长期依赖性的能力，而不产生额外的参数和计算负担？”

最近，状态空间模型（SSM）在研究人员中引起了相当大的关注。在经典SSM研究[10]奠定的基础上扩展，现代SSM（例如，Mamba [4]）不仅建立了远程依赖关系，而且还展示了关于输入大小的线性复杂性，使Mamba成为UNet轻量级道路上CNN和Transformer的强大竞争对手。一些当代的努力，如U-Mamba [14]，已经提出了一种混合CNN-SSM块，融合了局部特征提取能力与SSM在捕获纵向依赖关系方面的能力。然而，U-Mamba [14]引入了大量参数和计算负载（173.53 M参数和18，057.20 GFLOP），使得在移动的医疗保健设置中部署医疗分割任务具有挑战性。因此，在这项研究中，我们引入了LightM-UNet，这是一种基于Mamba的轻量级U形分割模型，它在实现最先进性能的同时显著降低了参数和计算成本（如图1所示）。这项工作的贡献有三个方面。

贡献：

（1）我们介绍LightM-UNet，它是UNet和Mamba的轻量级融合，仅拥有1 M的参数计数。通过对2D和3D真实数据集的验证，LightM-UNet超越了现有的最先进的模型。与著名的nnU-Net [8]和同期的UMAMBA [14]相比，LightM-UNet分别将参数数量减少了116倍和224倍。

（2）在技术上，我们提出了残余视觉曼巴层（RVM层），以纯粹的曼巴方式从图像中提取深层特征。以最小的引入新的参数和计算开销，我们进一步提高SSM的能力，通过利用残差连接和调整因子来增强视觉图像中的长距离空间依赖关系。

（3）坦率地说，与将UNet与Mamba集成的同期努力相反，我们主张将Mamba用作UNet内CNN和Transformer的轻量级替代品，旨在解决真实的医疗环境中计算资源限制带来的挑战。据我们所知，这代表了将Mamba引入UNet作为轻量级优化策略的开创性努力。

3、模型

RVM：

VSS Block：

解码器：

4、实验：

验证LightM-UNet的2D（2D X-ray images.）和3D（3D CT images）数据的性能。数据以7：1：2的比例随机分为训练集、验证集和测试集。

LightM-UNet使用PyTorch框架实现，并且三个编码器块中的RVM层的数量分别设置为1、2和2。所有实验均在单个Quadro RTX 8000 GPU上进行。SGD被用作优化器，初始化为1 e-4的学习速率。使用PolyLRScheduler作为调度器，总共训练了100个历元。此外，将损失函数设计为交叉熵损失和Dice损失的简单组合。对于LiTs数据集，将图像归一化并调整为128 × 128 × 128，批量为2。对于蒙哥马利和深圳数据集[9]，将图像归一化并调整大小为512 × 512，批量大小为12。

为了评估LightM-UNet，我们将其与医学图像分割竞赛中常用的两种基于CNN的分割网络（nnU-Net [8]和SegResNet [15]）、两种基于变换器的网络（UNETR [7]和SwinUNETR [6]）和一种基于曼巴的网络（U-Mamba [14]）进行了比较。此外，我们还采用了并集平均交集（mIoU）和骰子相似性评分（DSC）作为评价指标。

比较实验结果见表。1表明我们的LightM-UNet在LiTS数据集上实现了全面的最新性能[11]。值得注意的是，与nnU-Net等大型模型相比，LightM-UNet不仅表现出了上级的性能，而且显著地减少了47.39倍的参数数量和15.82倍的计算成本。与同时期的U-Mamba [14]相比，LightM-UNet在平均mIoU方面的性能提高了2.11%。特别是对于通常太小而不容易检测的肿瘤，LightM-UNet也能实现了3.63%的mIoU提升。重要的是，作为一种将Mamba融入UNet架构的方法，LightM-UNet与U-Mamba相比，仅使用了1.07%的参数和2.53%的计算资源。

2D数据集的实验结果总结在表中。2.LightM-UNet再一次实现了最佳性能，并显著超过了其他基于Transformer和Mamba的文献。此外，LightM-UNet的突出之处在于其显著的低参数计数，仅使用1.09M个参数。这表示与nnU-Net [8]和U-Mamba [14]相比，参数分别降低了99.14%和99.55%。为了更清楚地显示实验结果，请参见图1。图3所示的分割结果示例表明，与其他模型相比，LightM-UNet显示出更平滑的分割边缘，并且不会对小对象（如肿瘤）产生错误识别。