Source-linked AI summary
LightM-UNet: Mamba Assists in Lightweight UNet for Medical Image Segmentation
Weibin Liao, Yinghao Zhu, Xinyuan Wang, Chengwei Pan, Yasha Wang, Liantao Ma
TL;DR
Medical segmentation models need global context without the parameter and computational burden of Transformer-based approaches, especially for mobile healthcare. LightM-UNet replaces CNN and Transformer components within UNet with a lightweight Mamba-based design using RVM layers and long-range modeling modules. Across 2D and 3D datasets, it reports state-of-the-art segmentation with about 1M parameters and substantially lower resource costs than established baselines.
Problem
UNet-based segmentation is limited by convolutional locality, while Transformer-based approaches impose quadratic complexity and high computational costs for dense prediction.
Method
LightM-UNet integrates Mamba into a U-shaped network, using pure Mamba encoder structures, RVM layers, and modules for long-range spatial dependency modeling.
Results
LightM-UNet achieves state-of-the-art performance on 2D and 3D segmentation tasks with only 1M parameters, while reducing parameters by 116× versus nnU-Net.
Takeaways & Limitations
The results support using Mamba as a lightweight optimization strategy for UNet in resource-constrained medical segmentation settings.
Takeaways & Limitations
Future work includes designing a more lightweight network and validating it on more datasets covering multiple organs.
Abstract
from arXiv · showhide
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.
1 Introduction
UNet provides a strong encoder-decoder foundation for medical image segmentation, but convolutional locality and Transformer self-attention limit long-range modeling or impose high computational costs. LightM-UNet addresses these constraints by using Mamba as a lightweight UNet substitute and reports state-of-the-art performance with substantially reduced resource use.
- UNet’s convolutional locality limits explicit global and long-range semantic interactions, despite extensions using atrous convolutions, self-attention, and image pyramids.
- Transformer-based UNet solutions capture global information but incur quadratic complexity with image size, creating substantial overhead for dense medical segmentation.
- Mamba offers long-range dependency modeling with linear input-size complexity, making it a lightweight alternative to CNNs and Transformers within UNet.
- 173.53M parameters and 18,057.20 GFLOPs make U-Mamba challenging to deploy in mobile healthcare settings.
- LightM-UNet uses 1M parameters, surpasses existing state-of-the-art models on 2D and 3D datasets, and reduces parameters by 116× versus nnU-Net and 224× versus U-Mamba.
2 Methodologies
LightM-UNet uses a U-shaped encoder-decoder with depthwise convolutions for shallow features and resolution changes across encoder and decoder blocks. Pure Mamba encoder components, RVM layers, VSS modules, residual connections, and adjustment factors model long-range dependencies while keeping the design lightweight.
- Architecture Overview: LightM-UNet begins with depthwise convolution, then uses three encoder blocks that double channels and halve resolution at each stage.
- Architecture Overview: The bottleneck preserves feature-map size while modeling long-range spatial dependencies, and three decoder blocks halve channels and double resolution to restore the image.
- Encoder Block: Encoder blocks use solely Mamba structures, flattening feature maps before applying consecutive RVM Layers to extract deep features.
- RVM Layer: RVM Layers combine LayerNorm, VSSM processing, residual connections, and adjustment factors to enhance long-range spatial modeling with almost no added parameters or computational complexity.
- VSS Module: The VSS Module expands channels along two parallel branches, processes them with linear layers and nonlinear operations, combines them by Hadamard product, and projects back to the original channel count.
- Decoder Block: Decoder blocks add skip and previous-block features, then apply depthwise convolution, residual processing, ReLU, adjustment factors, and bilinear interpolation for resolution restoration.
3 Experiments
Experiments on LiTS and Montgomery&Shenzhen show that LightM-UNet achieves strong segmentation performance with substantially lower parameter and computation costs. Ablations support the value of long-range modeling and the RVM Layer components.
- Comparative Results: LightM-UNet achieves state-of-the-art performance on the LiTS 3D dataset while reducing parameters and computational costs versus nnU-Net by 47.39× and 15.82×.The comparison uses various 3D segmentation models.
- Comparative Results: 2.11% average mIoU improvement separates LightM-UNet from U-Mamba on LiTS, while tumor segmentation improves by 3.63%.The paper highlights tumors as small targets that are difficult to detect.
- Comparative Results: LightM-UNet achieves optimal performance on Montgomery&Shenzhen while using 1.09M parameters, 99.14% fewer than nnU-Net and 99.55% fewer than U-Mamba.The dataset evaluates the 2D version of the model.
- Comparative Results: Visual examples show smoother segmentation edges and no erroneous identification of small objects such as tumors compared with other models.Fig. 3 presents LiTS tumor and liver regions and Montgomery&Shenzhen lung regions, with arrows marking differing areas.
- Ablation Results: Transformer- and VSSM-based variants outperform convolution-based variants, supporting the benefit of modeling long-range dependencies.Replacing VSSM with convolution or self-attention causes performance sacrifices while adding parameters and computational overhead.
- Ablation Results: Removing RVM adjustment factors and residual connections reduces mIoU by 0.44% and 0.69% without meaningfully lowering parameter count or computational overhead.The ablation is conducted on LiTS, with additional Montgomery&Shenzhen analysis reported in supplementary material.
4 Conclusion
LightM-UNet is presented as a lightweight Mamba-based network for 2D and 3D medical image segmentation. The paper reports state-of-the-art performance with about 1M parameters and identifies broader multi-organ validation as future work.
- Conclusion: LightM-UNet achieves state-of-the-art performance in both 2D and 3D segmentation tasks with only 1M parameters and over 99% fewer parameters than recent Transformer-based architectures.The conclusion also reports significantly lower GFLOPs.
- Future Work: The authors identify validation on more datasets covering multiple organs as future work.They also aim to design a more lightweight network and foster applications in mobile health.