MHMamba: Multi-Head Mamba for 3D Brain Tumor Segmentation

技術概要：3D 脳腫瘍セグメンテーションのための MHMamba

問題定義
3D MRI における脳腫瘍セグメンテーションは、腫瘍形態とマルチモーダルコントラストの高度な不均一性により、重大な課題に直面している。畳み込みニューラルネットワーク（CNN）は局所的なテクスチャや境界のモデル化に優れているが、受容野が限られているため、長距離依存関係のモデル化には苦戦する。一方、Transformer ベースのモデルはグローバルな文脈を捉えることができるが、二次的な計算量とメモリ複雑性（$O(N^2)$）に悩まされており、 aggressive なダウンサンプリングやウィンドウ処理なしに高解像度の 3D ボリュームに適用することは困難である。これらは文脈の非一貫性や境界の不整合を引き起こす可能性がある。さらに、既存のハイブリッド手法では、畳み込みブランチとアテンションブランチ間の統計的ミスマッチ（例：バッチ/インスタンス正規化対レイヤー正規化）に起因する不安定な学習や特徴融合の問題をしばしば抱えている。

手法
本論文は、線形複雑性でグローバル特徴をモデル化しつつ高解像度の詳細を保持するように設計された U 字型エンコーダ - デコーダアーキテクチャである**MHMamba（Multi-Head Mamba）**を提案する。このフレームワークは、以下の 3 つの主要コンポーネントで構成される。

1. マルチヘッド Mamba（MHM）モジュール:

   • エンコーダ内の標準ブロックを、マルチヘッド状態空間モデル（SSM）に置き換える。
   • チャネル次元を並列な SSM ヘッド（$N_h$）に分割し、異なるヘッドが固定されたモダリティ割り当てなしに、マルチモーダル入力から相補的な長距離依存関係パターンを学習できるようにする。
   • 各ヘッドは 3D 空間特徴をシーケンスにシリアライズし、選択的状態更新（Mamba）を適用した後、出力を連結と射影を通じて融合する。
   • この設計は、マルチモーダル条件下でのグローバルモデル化能力と学習の安定性を高めつつ、線形複雑性を維持する。

2. チャネル - 空間較正アテンション（CSCA）モジュール:

   • マルチヘッド Mamba 処理後に特徴統計を整合させ、病変応答を強化するように設計されている。
   • 双経路メカニズムを採用する。
     • チャネル経路: グローバル平均プーリングと最大プーリングを用い、その後 MLP を通じてチャネル重みを生成する。
     • 空間経路: 4 つの空間統計（平均、標準偏差、最大値、最小値）を革新的に統合し、空間重みマップを生成して不均一性と顕著なピークを捉える。
   • ゲーティング機構がこれらのブランチを適応的に融合し、その後、元の特性応答を保持するために残差接続が追加される。

3. 適応的ゲーティング融合（AGF）モジュール:

   • エンコーダとデコーダ間のスキップ接続に位置する。
   • 高レベルのセマンティクスと低レベルの詳細間の対立に対処するため、マルチスケール特徴を動的に融合する。
   • 入力特徴はチャネル次元に沿ってサブグループに分割される。各サブグループに対して、モジュールはエンコーダ特徴とデコーダ特徴のバランスを取る独立した融合重みを学習し、境界予測の一貫性を欠く原因となる特徴の対立を防ぐ。

主な貢献

• マルチヘッド Mamba アーキテクチャ: 本論文は、チャネルを並列ヘッドに分解するマルチヘッド SSM 手法を 3D MRI に導入する。これにより、Mamba の線形複雑性の利点を維持しつつ、長距離モデル化とマルチモーダル適応性が強化される。
• 特徴較正と融合: 特徴スケールと統計を統一するCSCAモジュールの統合と、スキップ接続においてマルチスケール特徴を動的に融合するAGFモジュールの統合は、境界の一貫性と小容積病変の検出を大幅に向上させる。
• 体系的な検証: この手法はBraTS2021およびBraTS2023データセットで厳密に評価され、Dice スコアとハウスドルフ距離（HD95）の観点から、nnU-Net、SegMamba、Transformer ベースのモデルを含む強力なベースラインを上回る安定した改善を示した。

実験結果
BraTS2021 および BraTS2023 データセットで行われた実験により、MHMamba が最先端またはそれに準ずる性能を達成することが示された。

• BraTS2021: 平均 Dice スコア91.02%、平均 HD953.38を達成し、SegMamba（Dice 90.32%、HD95 4.18）や nnU-Net（Dice 89.40%、HD95 5.58）などのベースラインを上回った。
• BraTS2023: 平均 Dice スコア90.23%、平均 HD953.67を達成した。
• アブレーション研究:
  • CSCA、AGF、およびマルチヘッド Mamba モジュールを個別に追加することで、全腫瘍サブ領域（全腫瘍、腫瘍コア、増強腫瘍）において性能向上が確認された。
  • 最適な構成は4 ヘッドであり、表現能力と計算効率のバランスが取れていた。8 ヘッドに増やすと、3D 学習における小さなバッチサイズでの過分割と最適化ノイズにより、性能が低下した。
  • 定性的な結果は、MHMamba が特に不均一な腫瘍領域や小さな増強領域において、より断片化の少ない予測と緊密な境界を生成することを示している。

意義と主張
本論文は、MHMamba が 3D 脳腫瘍セグメンテーションにおけるグローバル文脈モデル化と局所詳細保持のトレードオフを効果的に解決すると主張している。マルチヘッドアーキテクチャを通じて状態空間モデルの線形複雑性を活用することで、この手法は Transformer の二次的なメモリオーバーヘッドを回避しつつ、CNN の局所受容野の限界を克服する。提案されたモジュール（CSCA および AGF）は、マルチモーダル 3D セグメンテーションでしばしば見られる統計的ミスマッチと文脈の非一貫性を特にターゲットとしており、放射線治療計画や術前評価などの臨床応用に適した、堅牢で効率的かつ正確なフレームワークを実現する。著者は、この手法が Mamba ベースのモデル化の効率性の利点を維持しつつ、腫瘍コアや小容積病変に対する感度を大幅に向上させていることを強調している。