SegMaFormer: 3D 医療画像セグメンテーションのための効率的なハイブリッド状態空間・トランスフォーマーモデル

本論文は、3D 医療画像セグメンテーションにおける計算コストと精度のバランスを最適化するため、SegMaFormerという軽量なハイブリッドアーキテクチャを提案しています。従来の CNN や大規模なトランスフォーマーモデルの課題を克服し、Mamba（状態空間モデル）とトランスフォーマーの長所を階層的に組み合わせることで、少ないパラメータ数と計算量（FLOPs）で最先端（SOTA）レベルの性能を実現しています。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細にまとめます。

---

1. 背景と問題定義

医療画像分析、特に 3D 体積データ（CT や MRI）のセグメンテーションにおいて、深層学習は重要な役割を果たしています。

• CNN の限界: 従来の畳み込みニューラルネットワーク（CNN）は局所的な特徴抽出に優れていますが、受容野が限られているため、画像全体の文脈（グローバルコンテキスト）を捉えるのが困難です。
• トランスフォーマーの課題: 自己注意機構（Self-Attention）を用いたトランスフォーマー（ViT, UNETR など）は長距離依存関係をモデル化できますが、3D 体積データをパッチに分割するとシーケンス長が爆発的に増加します。その結果、計算複雑度が $O(N^2)$ となり、メモリ消費と推論時間の面で実用的な課題（特にリソースが限られた医療現場）が生じます。
• Mamba の現状: 最近登場した状態空間モデル（SSM）、特に Mamba は線形計算量で長距離シーケンスを効率的に処理できますが、医療画像セグメンテーションにおいて単独で使用すると、CNN のような局所的なバイアスが不足し、精度が十分でない場合があります。

課題: 高い精度を維持しつつ、パラメータ数と計算量を大幅に削減できる、医療現場で実用的な軽量モデルの構築。

2. 提案手法：SegMaFormer

SegMaFormer は、Mamba の効率的なシーケンス混合能力と、トランスフォーマーの強力なグローバル文脈モデル化能力を、階層的なエンコーダ内でシナジーさせるハイブリッド構造を採用しています。

アーキテクチャの概要

• エンコーダ（4 段階の階層構造）:
  • 初期段階（高解像度）: 解像度が高くトークン数が多いため、計算コストの低いMamba ベースの層を使用します。これにより、空間的な文脈を効率的に捉えつつ、計算オーバーヘッドを最小化します。
  • 後期段階（低解像度）: 解像度が低下しトークン数が減るため、**自己注意機構（Self-Attention）**を使用します。トークン数が少ない段階では、$O(N^2)$ のコストが許容範囲となり、豊富なグローバル文脈のモデル化が可能になります。
• 3D 回転位置埋め込み（3D-RoPE）:
  • 従来の位置エンコーディングに加え、3D-RoPEを導入しました。これにより、ボクセル構造を維持しつつ、回転一貫性のある位置情報を Mamba とトランスフォーマーの両方に提供し、複雑な解剖学的領域における空間関係の捕捉能力を向上させています。
• デコーダ:
  • 重みなし（Convolution 不要）の軽量な MLP（多層パーセプトロン）ベースのデコーダを使用し、マルチスケールの特徴を融合させてセグメンテーションマスクを生成します。これにより、過剰なパラメータ化を回避しています。

計算効率化の戦略

• 3D 画像をパッチ化するとシーケンス長 $N$ が非常に大きくなるため、全解像度で自己注意を適用するのは非現実的です。
• 本手法は、高解像度段階では $O(N)$ の計算量を持つ Mamba を、低解像度段階でのみ $O(N^2)$ の自己注意を適用することで、全体として計算コストを劇的に削減しています。

3. 主要な貢献

1. ハイブリッドエンコーダの導入: 初期段階で Mamba ベースの状態空間層、後期段階で自己注意機構を組み合わせることで、3D 医療画像の長距離依存関係を効率的にモデル化しつつ、3D 注意の計算負荷を大幅に軽減しました。
2. 3D-RoPE の統合: 3D 重なりパッチ埋め込み後に 3D-RoPE を導入し、Mamba とトランスフォーマーの両方の性能を向上させ、解剖学的構造の空間的関係をより正確に捉えられるようにしました。
3. 軽量かつ高性能な実装:
   • パラメータ数：約 200 万（2M）
   • 計算量（FLOPs）：15.2 GFLOPs
   • これにより、既存の SOTA モデルと比較して最大75 倍のパラメータ削減と、著しい計算コストの低下を実現しました。

4. 実験結果

提案モデルは、3 つの主要な公開ベンチマーク（Synapse, BraTS, ACDC）で評価され、事前学習なしで競合モデルと同等以上の性能を示しました。

データセット	比較対象モデル	提案モデル (SegMaFormer) の結果
BraTS (脳腫瘍)	nnFormer (150.5M パラメータ)	平均 Dice 83.79% (nnFormer: 86.4%) ※パラメータ 75 倍以上少ないながら、実用的な精度を達成。
Synapse (腹部 CT)	U-Mamba (172.6M), nnFormer (150.5M)	平均 Dice 83.33% ※U-Mamba や nnFormer に次ぐ 3 位。巨大モデルに匹敵する精度を 2M パラメータで達成。
ACDC (心臓)	nnFormer, Primus-S	平均 Dice 91.11% ※SOTA モデル（nnFormer: 92.06%）と僅差（約 1% 未満）の性能。

• 特徴: 小器官や境界の明確な構造において、Mamba 単体モデルが局所性の欠如で性能が落ちる傾向がある中、3D-RoPE とハイブリッド構造によりこれを補い、CNN ベースのモデルや大規模トランスフォーマーと競合する精度を達成しました。

5. 意義と結論

• 医療現場への実用性: 計算リソースが限られた病院や発展途上国でも、高性能な 3D セグメンテーションを低コストのハードウェアで実行可能にします。
• アーキテクチャの指針: 「パラメータ数が多いほど良い」という従来の考え方を覆し、適切なアーキテクチャ設計（Mamba とトランスフォーマーの役割分担、位置エンコーディングの工夫）によって、軽量モデルでも SOTA 性能が達成可能であることを示しました。
• 今後の展望: 小器官や微細な解剖学的境界の認識精度はさらに向上の余地がありますが、本モデルはスケーラブルで臨床応用可能な 3D 医療画像セグメンテーションの基盤として確立されました。

総じて、SegMaFormer は、計算効率と高精度の両立という長年の課題に対し、Mamba とトランスフォーマーのハイブリッド化という革新的なアプローチで解決策を提示した画期的な研究です。