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

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

2024-03-11

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

Image and Video Processing,Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper primarily addresses the issue that existing U-Net and its variant models are not suitable for mobile medical applications due to their large number of parameters and heavy computational burden in medical image segmentation tasks. Specifically, traditional U-Net models are based on convolutional neural networks (CNNs) and have the inherent locality limitation of convolution operations, making it difficult for the model to capture global and long-range semantic information. Although some studies have attempted to enhance the model's ability to capture long-range dependencies by introducing attention mechanisms or transformer architectures, these methods often lead to a significant increase in computational complexity, especially in tasks requiring dense predictions such as medical image segmentation.

To solve the above problems, the authors propose a lightweight network structure—LightM-UNet, which combines the Mamba algorithm with the U-Net architecture. Mamba is a state space model (SSM) that can establish long-range dependencies with linear time complexity, making it very suitable for optimizing the computational efficiency of U-Net. LightM-UNet uses Mamba's Residual Vision Layer (RVM Layer) to extract deep semantic features and model long-range spatial dependencies while maintaining linear computational complexity.

The main contributions of the paper include:
1. Proposing LightM-UNet, a lightweight U-Net variant based on Mamba, containing only about 1 million parameters. Experimental results on 2D and 3D real-world datasets show that its performance surpasses the state-of-the-art while significantly reducing the number of parameters and computational cost.
2. Designing a technique called Residual Vision Mamba Layer (RVM Layer) to extract deep features from images in a pure Mamba manner, while using residual connections and adjustment factors to further enhance the SSM's ability to model long-range spatial dependencies.
3. For the first time, applying Mamba as an effective alternative to CNNs and transformers in U-Net, aiming to address the computational resource constraints in practical medical scenarios.

In summary, the main goal of this paper is to develop a medical image segmentation model that can effectively handle long-range dependencies while being lightweight, enabling deployment in resource-constrained environments such as mobile health applications.