[論文レビュー] VM-UNet: Vision Mamba UNet for Medical Image Segmentation

Jiacheng Ruan, Suncheng Xiang+1 more|arXiv (Cornell University)|Feb 4, 2024

Medical Image Segmentation Techniques被引用数 205

ひとこと要約

VM-UNetは、Vision Mambaブロックを使用する純粋な State Space Modelベースの医用画像分割のU-Netで、ISIC17/ISIC18およびSynapseデータセットで競合的な結果を達成します。

ABSTRACT

In the realm of medical image segmentation, both CNN-based and Transformer-based models have been extensively explored. However, CNNs exhibit limitations in long-range modeling capabilities, whereas Transformers are hampered by their quadratic computational complexity. Recently, State Space Models (SSMs), exemplified by Mamba, have emerged as a promising approach. They not only excel in modeling long-range interactions but also maintain a linear computational complexity. In this paper, leveraging state space models, we propose a U-shape architecture model for medical image segmentation, named Vision Mamba UNet (VM-UNet). Specifically, the Visual State Space (VSS) block is introduced as the foundation block to capture extensive contextual information, and an asymmetrical encoder-decoder structure is constructed with fewer convolution layers to save calculation cost. We conduct comprehensive experiments on the ISIC17, ISIC18, and Synapse datasets, and the results indicate that VM-UNet performs competitively in medical image segmentation tasks. To our best knowledge, this is the first medical image segmentation model constructed based on the pure SSM-based model. We aim to establish a baseline and provide valuable insights for the future development of more efficient and effective SSM-based segmentation systems. Our code is available at https://github.com/JCruan519/VM-UNet.

研究の動機と目的

• 医用画像分割における純粋なSSMベースモデルの探究を動機付ける。
• Vision Mambaブロック（VSS）を用いた非対称U-NetでVM-UNetアーキテクチャを提案する。
• 公開データセットでの純粋なSSMベース医用画像分割のベースラインを確立する。
• 皮膚病変および多臓器分割におけるVM-UNetを評価し競争力を検証する。

提案手法

• パッチ埋め込み/展開を用いた4段階の非対称エンコーダ-デコーダを利用する。
• エンコーダとデコーダの双方でコア特徴抽出器としてVision Mamba（VSS）ブロックを使用する。
• VSSブロックでは、長距離の文脈モデリングのためにSS2Dを用いた2つの分岐経路を適用する。
• SS2Dをスキャン展開/結合とMamba由来のS6ブロックで実装し、方向依存性を捉える。
• 加法融合による単純なスキップ接続を採用し、BceDiceまたはCeDice損失で学習する。
• VM-UNetをVMamba-S pretrained weightsで初期化し、ISIC17/ISIC18/Synapseデータセットで訓練する。

実験結果

リサーチクエスチョン

• RQ1 医用画像において純粋なSSMベースモデルが競争力のある分割性能を達成できるか？
• RQ2 Vision Mamba UNetは皮膚病変および臓器分割において、CNNベースおよびTransformerベースのベースラインとどう比較されるか？
• RQ3 事前学習済みのVMambaウェイトがVM-UNetの性能に与える影響は何か？
• RQ4 将来のSSMベース分割法におけるVM-UNetが設定するべきベースラインは何か？

主な発見

Dataset	Model	mIoU (%) ↑	DSC (%) ↑	Acc (%) ↑	Spe (%) ↑	Sen (%) ↑
ISIC17	UNet	76.98	86.99	95.65	97.43	86.82
ISIC17	UTNetV2	77.35	87.23	95.84	98.05	84.85
ISIC17	TransFuse	79.21	88.40	96.17	97.98	87.14
ISIC17	MALUNet	78.78	88.13	96.18	98.47	84.78
ISIC17	VM-UNet	80.23	89.03	96.29	97.58	89.90
ISIC18	UNet	77.86	87.55	94.05	96.69	85.86
ISIC18	UNet++	78.31	87.83	94.02	95.75	88.65
ISIC18	Att-UNet	78.43	87.91	94.13	96.23	87.60
ISIC18	UTNetV2	78.97	88.25	94.32	96.48	87.60
ISIC18	SANet	79.52	88.59	94.39	95.97	89.46
ISIC18	TransFuse	80.63	89.27	94.66	95.74	91.28
ISIC18	MALUNet	80.25	89.04	94.62	96.19	89.74
ISIC18	VM-UNet	81.35	89.71	94.91	96.13	91.12
Synapse	VM-UNet	DSC 81.08	HD95 19.21	Aorta 86.40	Gallbladder 69.41	Kidney(L) 86.16	Kidney(R) 82.76	Liver 94.17	Pancreas 58.80	Spleen 89.51	Stomach 81.40

• VM-UNetはISIC17およびISIC18で競合的なmIoU、DSC、および精度を達成し、いくつかのベースラインを上回る。
• ISIC17では、VM-UNetはmIoU 80.23%、DSC 89.03%、Acc 96.29%、Spe 97.58%、Sen 89.90%を達成。
• ISIC18では、VM-UNetはmIoU 81.35%、DSC 89.71%、Acc 94.91%、Spe 96.13%、Sen 91.12%を達成。
• SynapseではVM-UNetはDSC 81.08%およびHD95 19.21をデータセット全体で達成。
• Swin-UNet（純粋なTransformer）と比較して、DSCで1.95%、HD95で2.34mm上回る。
• アブレーションにより、VMamba-S pretrained weightsの使用がランダム初期化より顕著に性能を向上させることを示した。