[Paper Review] LightM-UNet: Mamba Assists in Lightweight UNet for Medical Image Segmentation

Weibin Liao, Yinghao Zhu+4 more|arXiv (Cornell University)|Mar 8, 2024

COVID-19 diagnosis using AI46 citations

TL;DR

LightM-UNet combines Mamba with UNet to create a lightweight 3D/2D medical image segmentation model with ~1M params, achieving state-of-the-art results on LiTS while greatly reducing parameters and FLOPs.

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.

Motivation & Objective

• Motivate lightweight medical image segmentation suitable for mobile health by reducing computational burden and parameters.
• Leverage Mamba (State Space Model) to capture long-range dependencies within UNet without Transformer-level costs.
• Propose Residual Vision Mamba Layer (RVM Layer) and Vision State-Space Module (VSS Module) to enhance deep feature modeling with minimal parameter increase.
• Demonstrate state-of-the-art performance on 3D LiTS and 2D Montgomery&Shenzhen datasets while keeping parameter count around 1M.

Proposed method

• Embed Mamba as a pure Mamba-based encoder within a UNet-like architecture.
• Use Residual Vision Mamba Layer (RVM Layer) with LayerNorm, VSSM, and residual adjustment factors to model long-range dependencies.
• Incorporate Vision State-Space Module (VSS Module) with parallel branches and Hadamard product to fuse local and global features.
• Structure LightM-UNet with 3 Encoder Blocks, a Bottleneck Block, and 3 Decoder Blocks in a UNet-style encoder–bottleneck–decoder layout.
• Apply depthwise convolutions for shallow feature extraction and bilinear upsampling at the end to produce segmentation maps.
• Train with a combination of Cross-Entropy and Dice losses on 3D LiTS and 2D Montgomery&Shenzhen datasets, comparing against nnU-Net, SegResNet, UNETR, SwinUNETR, and U-Mamba.

Experimental results

Research questions

• RQ1 Can UNet-based medical image segmentation achieve global long-range modeling with a lightweight architecture?
• RQ2 How does a pure Mamba-based encoder compare to CNN/Transformer blocks in terms of accuracy and efficiency for 3D/2D segmentation?
• RQ3 Do RVM Layer and VSS Module provide meaningful gains with negligible parameter overhead?
• RQ4 Is LightM-UNet competitive against state-of-the-art models while drastically reducing parameters and FLOPs?

Key findings

• LightM-UNet achieves 1.87M parameters and 457.62 GFLOPs for the 3D LiTS task.
• On LiTS, LightM-UNet attains Liver DSC 96.31 and Tumor DSC 72.86 with average DSC 84.58 and average mIoU 77.48.
• Compared to nnU-Net, LightM-UNet reduces parameters and FLOPs by 47.39x and 15.82x respectively while achieving higher average DSC/mIoU.
• Against U-Mamba, LightM-UNet improves average mIoU by 2.11 percentage points and tumor DSC by 3.63 percentage points.
• Ablation shows replacing VSSM with CNN or Self-Attention degrades performance and increases parameters, while removing Adjustment factors or Residual connections lowers performance without saving cost.
• Overall, LightM-UNet delivers state-of-the-art performance with an ultra-lightweight footprint on 3D LiTS and competitive results on 2D datasets.