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

TL;DR

LightM-UNet introduces a lightweight medical image segmentation model using Mamba, achieving superior performance with significantly fewer parameters and computational costs, suitable for mobile health applications.

Contribution

This paper presents LightM-UNet, a novel lightweight segmentation model integrating Mamba into UNet, reducing complexity while improving accuracy over existing methods.

Findings

Outperforms state-of-the-art models on real-world datasets.

Reduces parameters by 116 times and computation by 21 times compared to nnU-Net.

Achieves better segmentation performance with lower resource requirements.

Abstract

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.