DAMamba-UNet3D: A Parameter-Efficient Mamba State Space U-Net with Dynamic Adaptive Scan for 3D Medical Image Segmentation

技術要約: DAMamba-UNet3D

問題提起
3D医療画像セグメンテーションには、厳格な計算およびパラメータ予算を遵守しつつ、精密な局所境界の描写と一貫したグローバルな空間コンテキストのバランスを取ることが求められる。既存のアーキテクチャには明確なトレードオフが存在する：

• 畳み込みU-Net: パラメータ効率が高く、$O(n)$ の局所混合を実現するが、明示的なグローバルコンテキストに欠け、長距離依存性を捉えるために深さのスケールアップが必要となる。
• ビジョンTransformer: 自己注意機構（Self-attention）を通じてグローバルな推論を提供するが、大規模な3Dボリューム（例：$128^3$ のパッチは約210万トークンを含む）に対しては、法外な $O(n^2)$ のメモリおよび計算コストが発生する。
• 状態空間モデル (SSM/Mamba): 情報を $O(n)$ の時間およびメモリでグローバルに伝播させる。しかし、既存の医療用SSMセグメンター（例：SegMamba）は、固定されたスキャンパターンと大きなパラメータ予算（約70Mパラメータ）に依存しており、データ依存の適応性に欠けている。

Dynamic Adaptive Scan (DAS) は、選択的スキャンの前にデータ依存の並べ替えを学習するために2D自然画像認識において導入されているが、医療画像への適用や3Dボリュームへの拡張は行われていない。

手法
著者らは、Tri-Plane 3D-DAS ブロックを U-Net フレームワークに統合したハイブリッド・エンコーダ・デコーダ・アーキテクチャである DAMamba-UNet3D を提案する。

1. アーキテクチャ設計:

   • ハイブリッド構造: 局所的な詳細を保持するために、ステム（E1）、ボトルネック、およびデコーダ全体に畳み込みブロックを維持する。
   • エンコーダの統合: Tri-plane DAMamba ブロックを、エンコーダのステージ E2–E4 にのみ挿入する。
   • コンパクト・バリアント: 主要モデルはベース幅 $C=16$ を使用し、約 530万パラメータ（SegMambaより約13倍少ない）となる。
   • 大規模バリアント (DAMamba-L): エンコーダに DAMamba ブロックをスタックし、畳み込みボトルネックを備えた、より広いバリアント（約70Mパラメータ）。直接比較のために SegMamba の容量に合わせるよう設計されている。

2. Tri-Plane Dynamic Adaptive Scan (DAS):

   • メカニズム: 固定スキャン（行優先、ジグザグ、または多方向の平坦化）とは異なり、DAS は入力依存の空間的な並べ替えを学習する。オフセット予測ネットワーク (OPN) が特徴マップから2Dオフセットを予測し、それらが選択的スキャンの前にバイリニア補間を介して特徴をリサンプリングするために使用される。
   • 3Dへの拡張: 2D DAS オペレータを、3つの解剖学的平面（Axial: $H \times W$、Coronal: $D \times W$、Sagittal: $D \times H$）にわたって適用することでボリュームへと拡張する。出力は平均化され、3D DAS ブロックを形成する：
     $$DAS3D(X) = \frac{1}{3}(DAS_{ax} + DAS_{cor} + DAS_{sag})$$
   • 選択的スキャン: 並べ替えられた特徴マップは標準的な Mamba 選択的スキャンブロックによって処理され、二次的なアテンションマップを伴わずに $O(N)$ のコストでグローバルなコンテキスト伝播を可能にする。

3. 学習および評価:

   • データセット: BraTS 2020 トレーニングセット（369症例、4モダリティ）。
   • 損失関数: 困難なボクセルや境界に焦点を当てるための Modulated Combo Loss ($L = (L_{CE} + \delta) \cdot (L_{mod}^{Dice} + \delta)$)。
   • プロトコル: 5分割被験者レベル交差検証、およびフルボリューム推論（スライディングウィンドウ方式、$128^3$ パッチ）。指標には、Enhancing Tumor (ET)、Tumor Core (TC)、および Whole Tumor (WT) に対する平均 Dice および HD95 を含む。

主な貢献

1. DASの初の医療応用: Dynamic Adaptive Scan を、2D自然画像分類から高密度3D医療セグメンテーションへと転用した。
2. 2Dから3Dへの拡張: 手設計された多方向平坦化の学習可能な代替案として、トライプレーン適応型スキャニングを介して DAS オペレータをボリュームデータへと引き上げた。
3. パラメータ効率: エンコーダにのみ Tri-plane DAS ブロックを配置したハイブリッド U-Net により、より大きなモデルと同等の性能を、約13倍低いパラメータコストで達成した。
4. 2スケールでの検証: コンパクトスケール（$\sim 5.3$M）と大規模スケール（$\sim 70$M）の両方でアプローチを検証し、アーキテクチャの配置が最適化されている場合、学習された DAS が固定スキャンパターンよりも優れた性能を示すことを実証した。

結果

• コンパクトモデル (DAMamba-UNet3D, $\sim 5.3$M): BraTS 2020 において平均 Dice 0.815 $\pm$ 0.013 を達成。これは、約13倍少ないパラメータ数であるにもかかわらず、再学習された SegMamba (0.824 $\pm$ 0.014) と同等の性能である。
• 大規模モデル (DAMamba-L, $\sim 70$M): 再学習された SegMamba を上回り、平均 Dice 0.829 $\pm$ 0.012 を達成（+0.5 ポイントの向上）。最大の利得は Tumor Core (TC) クラス（0.842 vs 0.827）で観察された。
• アブレーション研究:
  • 配置: DAS ブロックをエンコーダ（E2–E4）のみに配置した場合が最良の結果（0.815）をもたらした。ボトルネックまたはデコーダに DAS を追加すると性能が低下した（それぞれ 0.807 および 0.805）。これは、エンコーダ限定の配置が極めて重要であることを示している。
  • スキャンパターン: 学習された Tri-plane DAS は、エンコーダの配置が一致している場合、固定スキャン順序（Raster, Snake, Cross, Tri-orientated）よりも優れた性能を示した。

意義および主張
本論文は、ハイブリッド U-Net における学習された Tri-plane DAS は、3D 脳腫瘍セグメンテーションにおける固定された Tri-orientated Mamba (ToM) スキャニングに対する競争力のある代替案であると主張している。

• 効率性: コンパクトモデルは、ブルートフォース的な容量拡大ではなく、適応型スキャニングを活用することで、高い性能が大幅に削減されたパラメータ数で達成できることを示している。
• 性能: 大規模スケールにおいて、学習された適応メカニズム (DAS) は、同一の訓練および評価プロトコルの下で、固定スキャン軌跡 (ToM) を改善することができる。
• 謙虚な姿勢: 著者らは、BraTS 2020 における結果は有望であるが、本研究は単一のデータセットに限定されていることに注意を促している。彼らは、「DAS と ToM に関する一般的な主張を行う前に、マルチデータセットによる検証と統計的テストが必要である」と述べている。主要な貢献は、DAS を 3D 医療ボリュームへと拡張し、そのハイブリッド・パラメータ効率的アーキテクチャにおける生存能力を実証したことにある。