[論文レビュー] U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation

Jun Ma, Li Fei-Fei+1 more|arXiv (Cornell University)|Jan 9, 2024

AI in cancer detection被引用数 213

ひとこと要約

U-Mambaは、長距離依存性を線形スケーリングで捉え、データセット間で自己構成可能なハイブリッドCNN-SSMネットワークで、生体医用画像分割においてCNN-およびTransformerベースのライバルよりも複数タスクで優れている。

ABSTRACT

Convolutional Neural Networks (CNNs) and Transformers have been the most popular architectures for biomedical image segmentation, but both of them have limited ability to handle long-range dependencies because of inherent locality or computational complexity. To address this challenge, we introduce U-Mamba, a general-purpose network for biomedical image segmentation. Inspired by the State Space Sequence Models (SSMs), a new family of deep sequence models known for their strong capability in handling long sequences, we design a hybrid CNN-SSM block that integrates the local feature extraction power of convolutional layers with the abilities of SSMs for capturing the long-range dependency. Moreover, U-Mamba enjoys a self-configuring mechanism, allowing it to automatically adapt to various datasets without manual intervention. We conduct extensive experiments on four diverse tasks, including the 3D abdominal organ segmentation in CT and MR images, instrument segmentation in endoscopy images, and cell segmentation in microscopy images. The results reveal that U-Mamba outperforms state-of-the-art CNN-based and Transformer-based segmentation networks across all tasks. This opens new avenues for efficient long-range dependency modeling in biomedical image analysis. The code, models, and data are publicly available at https://wanglab.ai/u-mamba.html.

研究の動機と目的

• 生体医用画像分割における長距離依存性のモデリング課題に対処する。
• 局所特徴抽出と長距離コンテキストを組み合わせた汎用CNN-SSMバックボーン（U-Mamba）を提案する。
• データセットの多様性に自動適応する自己構成機構を組み込む。
• 4つの異なる生体医用分割タスク（3D腹部CT/MRI、内視鏡機器、顕微鏡細胞）でU-Mambaを評価し、CNNおよびTransformerベースのベースラインと比較する。
• 再現性とさらなる研究を促進するため、オープンソースのコード、モデル、およびデータを提供する。

提案手法

• 長距離依存性を線形スケーリングと入力依存選択で捉えるSelective Structured State Space Sequence Models (SSMs)をScan (S6)と組み合わせたMambaを導入する。
• U-Net様のエンコーダ-デコーダ（U-Mamba）に2つのResidualブロックをMambaブロックの前に配置して、局所情報とグローバル情報を融合する。
• 3D/2Dの特徴マップをMambaブロックへ適したシーケンス形状にフラット化し、並列ブランチのHadamard積融合を適用する。
• 2つのネットワーク変種を提供する：U-Mamba_Bot（ボトルネックのみMamba）とU-Mamba_Enc（全エンコーダブロックにMamba）。
• nnU-Netの自己構成フレームワークと統合して、データセットごとにアーキテクチャの深さとパッチサイズを自動的に適応する。
• 公平な比較を確保するため、固定のトレーニングプロトコルでDiceとクロスエントロピーロスを組み合わせて学習する。

実験結果

リサーチクエスチョン

• RQ1 U-MambaはCNNとTransformerを超える長距離依存性モデリングを生体医用画像分割で向上させることができるか。
• RQ2 線形スケーリングのMambaブロックは、多様な2D/3Dの生体医用分割タスクで実用的な利得をもたらすか。
• RQ3 腹部臓器、機器、細胞のセグメンテーションに対して、強力なCNN-およびTransformerベースのベースラインと比較してどうか。
• RQ4 U-Mambaの自己構成機構（nnU-Netを介して）は、手動チューニングなしでデータセットを跨いで堅牢な適応を可能にするか。

主な発見

Methods	Organs in Abdomen CT DSC	Organs in Abdomen CT NSD	Organs in Abdomen MRI DSC	Organs in Abdomen MRI NSD
nnU-Net	0.8615 ± 0.0790	0.8972 ± 0.0824	0.8309 ± 0.0769	0.8996 ± 0.0729
SegResNet	0.7927 ± 0.1162	0.8257 ± 0.1194	0.8146 ± 0.0959	0.8841 ± 0.0917
UNETR	0.6824 ± 0.1506	0.7004 ± 0.1577	0.6867 ± 0.1488	0.7440 ± 0.1627
SwinUNETR	0.7594 ± 0.1095	0.7663 ± 0.1190	0.7565 ± 0.1394	0.8218 ± 0.1409
U-Mamba_Bot	0.8683 ± 0.0808	0.9049 ± 0.0821	0.8453 ± 0.0673	0.9121 ± 0.0634
U-Mamba_Enc	0.8638 ± 0.0908	0.8980 ± 0.0921	0.8501 ± 0.0732	0.9171 ± 0.0689

• U-Mamba_Botは腹部CTでDSC 0.8683、NSD 0.9049を達成し、CNNとTransformerベースのベースラインを上回る。
• U-Mamba_EncはCTでDSC 0.8638、NSD 0.8980を達成し、腹部MRIでDSC 0.8501、NSD 0.9171を示し、クロスモダリティ性能が高い。
• 2Dタスクでは、MRIの臓器、内視鏡の機器、顕微鏡の細胞において、ベースラインモデルより平均性能が優れている。
• データセット全体で、U-Mambaは分割外れ値が少なく、肝臓・胆嚢などの難解な臓器や機器の境界の描出が改善されている。