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

UNet及其变体在医学图像分割中被广泛应用。然而，这些模型，特别是基于Transformer架构的模型，由于参数数量和计算负荷巨大，使它们不适用于移动医疗应用。最近，状态空间模型（SSM），例如Mamba，已经成为CNN和Transformer架构的竞争对手。基于此，我们将Mamba作为CNN和Transformer的轻量级替代品，应用于UNet中，旨在解决实际医疗环境中的计算资源限制所带来的挑战。为此，我们引入了轻量级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和其变种模型（如基于Transformer架构的模型）存在的参数数量和计算负担过大的问题，针对真实医疗环境中的计算资源限制，提出一种轻量级的解决方案。

论文提出了一种轻量级的解决方案——LightM-UNet，该方案将Mamba作为CNN和Transformer的轻量级替代品，将Mamba与UNet集成在一个轻量级框架中。LightM-UNet利用Residual Vision Mamba Layer以纯Mamba方式提取深度语义特征和建模长程空间依赖关系，具有线性计算复杂度。

实验结果表明，与现有的最先进文献相比，LightM-UNet在两个真实的2D/3D数据集上实现了超越，并且大幅度减少了参数和计算成本，分别为116倍和21倍。这凸显了Mamba在促进模型轻量化方面的潜力。作者已经公开了代码实现。

最近的相关研究包括nnU-Net等。