[논문 리뷰] LightM-UNet: Mamba Assists in Lightweight UNet for Medical Image Segmentation

Weibin Liao, Yinghao Zhu외 4명|arXiv (Cornell University)|2024. 03. 08.

COVID-19 diagnosis using AI인용 수 46

한 줄 요약

LightM-UNet은 Mamba를 UNet과 결합하여 약 1M 매개변수를 갖는 경량 3D/2D 의료 영상 분할 모델을 생성하고 LiTS에서 최첨단 성능을 달성하는 동시에 매개변수 수와 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.

연구 동기 및 목표

• 모바일 헬스에 적합한 경량 의료 영상 분할을 동기화하기 위해 계산 부담과 매개변수를 감소시킨다.
• 변환기 수준의 비용 없이 UNet 내에서 긴 범위 의존성을 포착하기 위해 Mamba (State Space Model)을 활용한다.
• 최소한의 매개변수 증가로 심층 특징 모델링을 강화하기 위해 Residual Vision Mamba Layer (RVM Layer) 및 Vision State-Space Module (VSS Module)을 제안한다.
• 약 1M 근처의 매개변수 수를 유지하면서 3D LiTS 및 2D Montgomery&Shenzhen 데이터셋에서 최첨단 성능을 입증한다.

제안 방법

• UNet 유사 구조 내에서 순수 Mamba 기반 인코더로 Mamba를 임베드한다.
• LayerNorm, VSSM, 잔여 보정 인자를 갖춘 Residual Vision Mamba Layer (RVM Layer)로 긴 범위 의존성을 모델링한다.
• 로컬 및 글로벌 특징 융합을 위해 병렬 분기와 Hadamard 곱을 갖춘 Vision State-Space Module (VSS Module)을 도입한다.
• UNet 스타일의 인코더–병목–디코더 배치에서 3개의 인코더 블록, 하나의 병목 블록, 3개의 디코더 블록으로 LightM-UNet를 구성한다.
• 얕은 특징 추출을 위한 깊이별 합성곱(depthwise convolutions)과 끝에서의 양선형 업샘플링을 적용하여 분할 맵을 생성한다.
• 3D LiTS 및 2D Montgomery&Shenzhen 데이터셋에서 Cross-Entropy 및 Dice 손실의 조합으로 학습하고 nnU-Net, SegResNet, UNETR, SwinUNETR, U-Mamba와 비교한다.

실험 결과

연구 질문

• RQ1 경량 아키텍처로 UNet 기반 의료 영상 분할이 글로벌 긴 범위 모델링을 달성할 수 있는가?
• RQ2 순수 Mamba 기반 인코더가 3D/2D 분할의 정확도와 효율성 측면에서 CNN/Transformer 블록과 어떻게 비교되는가?
• RQ3 RVM Layer와 VSS Module이 매개변수 증가 없이 의미 있는 이점을 제공하는가?
• RQ4 매개변수와 FLOPs를 대폭 줄이는 한편 LightM-UNet가 최첨단 모델과 경쟁력 있는가?

주요 결과

모델	파라미터(M)	GFLOPs	간 DSC	간 mIoU	종양 DSC	종양 mIoU	평균 DSC	평균 mIoU
LightM-UNet	1.87	457.62	96.31	92.92	72.86	62.05	84.58	77.48

• LightM-UNet은 3D LiTS 작업에서 1.87M 매개변수와 457.62 GFLOPs를 달성한다.
• LiTS에서 LightM-UNet은 Liver DSC 96.31 및 Tumor DSC 72.86, 평균 DSC 84.58 및 평균 mIoU 77.48을 달성한다.
• nnU-Net과 비교하여 LightM-UNet은 매개변수와 FLOPs를 각각 47.39배와 15.82배 감소시키면서 더 높은 평균 DSC/mIoU를 달성한다.
• U-Mamba에 비해 LightM-UNet은 평균 mIoU를 2.11퍼센트 포인트, 종양 DSC를 3.63퍼센트 포인트 개선한다.
• 제거 실험은 VSSM을 CNN 또는 Self-Attention으로 대체하면 성능이 저하되고 매개변수도 증가하며, 보정 인자나 잔여 연결을 제거하면 비용 절감 없이 성능이 저하된다.
• 전반적으로 LightM-UNet은 3D LiTS에서 최첨단 성능과 초경량 발자국, 2D 데이터셋에서 경쟁력 있는 결과를 제공합니다.