[論文レビュー] 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 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.

研究の動機と目的

• 計算量とパラメータを削減することで、モバイルヘルスに適した軽量な医用画像分割を推進する。
• Mamba（State Space Model）を活用して、Transformerレベルのコストを要せずにUNet内の長距離依存性を捉える。
• 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つのエンコーダーブロック、1つのボトルネックブロック、3つのデコーダーブロックを備えるよう LightM-UNet を構成する。
• 浅い特徴抽出には深さ方向畳み込みを適用し、最終的には双一次補間でアップサンプリングしてセグメンテーションマップを生成する。
• 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 LightM-UNetは、パラメータとFLOPsを劇的に削減しつつ、最先端モデルと競合できるか？

主な発見

モデル	パラメータ(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データセットでは競争力のある結果を示す。