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.
翻译: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在促进模型轻量化方面的潜力。我们的代码实现已在https://github.com/MrBlankness/LightM-UNet公开提供。