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

Jinhong Wang, Jintai Chen+2 more|arXiv (Cornell University)|Mar 12, 2024

Medical Image Segmentation Techniques被引用数 4

ひとこと要約

本稿では、大規模な窓サイズのS状態空間モデル（SSM）を用いて、線形計算量のMambaの特性を活かし、大規模な受容 field を実現する、2次元および3次元医療画像セグメンテーション向けの大規模窓ベースのMamba U-Net（LMa-UNet）を提案する。階層的かつ双方向のMambaブロック（LMブロック）を導入し、ピクセル単位およびパッチ単位のSSMを組み合わせることで、優れた局所的およびグローバル特徴モデリングを実現し、マルチオルガニズムセグメンテーションベンチマークにおいて、CNN、Transformer、および先行のMambaベースのモデルを上回るDSCおよびNSDスコアを達成した。

ABSTRACT

In clinical practice, medical image segmentation provides useful information on the contours and dimensions of target organs or tissues, facilitating improved diagnosis, analysis, and treatment. In the past few years, convolutional neural networks (CNNs) and Transformers have dominated this area, but they still suffer from either limited receptive fields or costly long-range modeling. Mamba, a State Space Sequence Model (SSM), recently emerged as a promising paradigm for long-range dependency modeling with linear complexity. In this paper, we introduce a Large Kernel Vision Mamba U-shape Network, or LKM-UNet, for medical image segmentation. A distinguishing feature of our LKM-UNet is its utilization of large Mamba kernels, excelling in locally spatial modeling compared to small kernel-based CNNs and Transformers, while maintaining superior efficiency in global modeling compared to self-attention with quadratic complexity. Additionally, we design a novel hierarchical and bidirectional Mamba block to further enhance Mamba's global and neighborhood spatial modeling capability for vision inputs. Comprehensive experiments demonstrate the feasibility and the effectiveness of using large-size Mamba kernels to achieve large receptive fields. Codes are available at https://github.com/wjh892521292/LKM-UNet.

研究の動機と目的

• CNNおよびTransformerの医療画像セグメンテーションにおける制限、特に受容 field の制限および長距離モデリングのための高コストを解決すること。
• 構造的状態空間モデル（SSM）としてのMambaの可能性を検証し、線形計算量で大規模な受容 field を実現する空間モデリングを可能にすること。
• 視覚タスクにおけるMambaの単方向性および位置情報の欠如を克服するため、双方向的かつ階層的なSSMアーキテクチャを設計すること。
• 大規模窓ベースのMambaモジュールが、2次元および3次元医療画像における局所的およびグローバル依存関係を効果的にモデリングできることを示すこと。
• マルチオルガニズムセグメンテーションデータセットにおける包括的なアブレーションおよびベンチマークを通じて、提案されたLMa-UNetの有効性を検証すること。

提案手法

• 標準的な畳み込みおよび自己注意モジュールを大規模窓SSMに置き換えることで、大規模な受容 field を実現する大規模窓ベースのMamba U-Net（LMa-UNet）を提案する。
• ピクセル単位のSSM（PiM）による局所的近傍モデリングと、パッチ単位のSSM（PaM）によるグローバル依存関係学習を組み合わせた、新しい階層的かつ双方向のMambaブロック（LMブロック）を導入する。
• トークナイゼーション後の空間的不連続性が引き起こすSSMの忘却問題を緩和し、局所的特徴表現を強化するためにPiMを採用する。
• 自己注意の二次的コストを回避し、線形計算量でパッチ間の長距離依存関係をモデリングするためにPaMを用いる。
• シーケンスモデリングにおける入力順序バイアスを低減し、位置情報の認識を向上させるために、双方向Mamba（BiM）を適用する。
• PyTorchを用いてnnU-Netフレームワークに実装し、各段階で適応的な窓サイズを採用し、標準的な訓練ハイパーパramータ（Adam、重み減衰、学習率）を用いる。

実験結果

リサーチクエスチョン

• RQ1 Mambaベースのモデルが大規模窓サイズを採用することで、小カーネルCNNおよび窓制限付きTransformerと比較して、医療画像セグメンテーションにおける優れた局所的およびグローバル特徴モデリングを達成できるか？
• RQ2 LMブロックにおけるPiMおよびPaMの階層的設計が、視覚タスクにおけるMambaの表現能力をどのように向上させるか？
• RQ3 双方向Mambaが、単方向SSMに内在する位置バイアスおよびシーケンス順序感受性をどの程度軽減できるか？
• RQ4 Mambaベースのモデルにおける窓サイズの拡大が、医療画像セグメンテーションにおいて測定可能な性能向上をもたらすか？
• RQ5 PiM、PaM、BiMの各コンponentが、LMa-UNet全体の性能に果たす相対的寄与度はどの程度か？

主な発見

• Abdomen CTデータセットにおいて、LMa-UNetはDSC 86.82およびNSD 90.02を達成し、CNNベース、Transformerベース、および先行のMambaベースのモデルを上回った。
• Abdomen MRデータセットでは、DSC 77.35およびNSD 83.80を達成し、評価されたすべてのベンチマークで最先端の性能を示した。
• 窓サイズを[10,5,5,5,5,5,5]から[40,20,20,10,10,5,5]に増加させることで、DSCは75.89から77.35に、NSDは82.26から83.80に向上し、大規模な受容 field の利点を裏付けた。
• アブレーションスタディの結果、PiMがPaMよりも性能向上に寄与していることが示され、局所的受容 field の強化が、グローバルモデリングそのものよりも重要であることがわかった。
• PiM + PaM + BiMのすべてのコンponentsを備えた完全なLMa-UNetは、DSC 77.35およびNSD 83.80を達成し、階層的かつ双方向設計の有効性を検証した。
• Mambaの線形計算量のおかげで、効率性を維持しつつ、従来のCNNおよびTransformerでは不可能な大規模窓モデリングを実現した。