Brain-SAM: A SAM-based Model Tailored for Brain MRI Lesion Segmentation

Brain-SAM に関する技術要約

問題提起
磁気共鳴画像法（MRI）は神経画像解析の基本であるが、脳構造や病変の正確なセグメンテーションは依然として困難である。Segment Anything Model (SAM) やその後継である SAM2 のような基盤モデルは、自然画像において驚異的な汎用性を示しているが、それらを脳MRIに適用するには大きな障壁が存在する。既存の医学的適応モデル（例：MedSAM、SAM-Med3D）は、心臓、肝臓、肺などの構造を含む混合臓器データセットで訓練されているのが一般的である。これらのモデルは、複雑な組織構成、不明瞭な灰白質・白質の境界、高度に多様な病変形態といった、脳特有の解剖学的特性への適応に苦慮している。さらに、ほとんどの SAM ベースの適応手法は、3D ボリュームスキャンを 2D スライスへと変換しており、スライス間の空間的コンテキストを無視している。加えて、手動によるインタラクティブなプロンプトへの依存は、臨床医がてんかん病変のような微細な病理に対して正確なプロンプトを提供することに困難を感じる可能性があるため、スケーラビリティを制限する。最後に、明確な境界を持つ 2D 自然画像向けに設計されたオリジナルの SAM マスクデコーダーは、境界がぼやけ、組織の不均一性が高い脳 MRI データに対しては、最適ではない性能を示すことが多い。

手法
著者らは、3D 脳 MRI 病変セグメンテーションに特化したモデルである Brain-SAM を提案する。このアーキテクチャは、主に 3 つの修正を加えることで SAM2 フレームワークを拡張している。

1. 画像のエンコーダーの 3D 適応: オリジナルの 2D Hiera エンコーダー（階層型 Vision Transformer）を、3D ボリュームデータを直接処理できるように適応させている。これには、以下の 4 つの主要モジュールの拡張が含まれる。

   • パッチ埋め込み (Patch Embedding): 2D 畳み込みを 3D 畳み込みに置き換え、ボリューム入力をボクセルパッチに分割する。
   • 位置埋め込み (Position Embedding): 深さ、高さ、幅にわたる空間的関係を捉えるために、2D から 3D 表象へと拡張する。
   • トークンシーケンス (Token Sequence): ボリューム的なアテンションに適した長いトークンシーケンスを作成するために、3D ボリュームをすべての次元に沿ってフラット化する。
   • Q-Pooling: MaxPool2d を MaxPool3d に置き換え、深さおよび空間次元にわたる特徴量のダウンサンプリングを可能にする。

2. UNETR に着想を得たデコーダー: 微細な局所テクスチャや境界の詳細を扱う際のオリジナル SAM デコーダーの限界に対処するため、著者らは UNETR に着想を得たデコーダーを導入している。このコンポーネントは、マルチスケール特徴融合を介して、Transformer のグローバルモデリング能力と CNN の局所的な微細表現を統合する。解像度の変換を必要とする標準的な UNETR デコーダーとは異なり、この設計は Hiera エンコーダー固有の階層的特徴を活用して、直接的な融合と漸進的なアップサンプリングを実現する。このアーキテクチャは、完全自動セグメンテーション用と、プロンプトベースのセグメンテーション用の中間特徴生成用の 2 つのデュアルヘッド出力をサポートする。

3. 学習およびファインチューニング戦略:

   • 事前学習: モデルは SAM2 の重みで初期化され、堅牢な医学的表現を確立するために、BraTS2021 データセット（1,000 ケース、4 モダリティ）を用いて完全自動セグメンテーションに特化した事前学習段階を経る。
   • アダプターベースのファインチューニング: キャタストロフィック忘却を防ぎつつ、複数の脳病変データセット（腫瘍、脳卒中、てんかん）に適応するために、各 Hiera ブロックの前に軽量なアダプターモジュール（3D 畳み込みを備えたボトルネック構造）を挿入する。ファインチューニング中、バックボーンは凍結されたまま、共有されるアダプターパラメータのみが更新される。
   • プロンプトメカニズム: ポイントプロンプトを主要なインタラクション手法として使用し、これを 3D へと拡張する。マスク生成戦略は、ドット積演算において更新されたキーを使用する代わりに、UNETR に着想を得たデコーダーからの中間特徴表現と処理されたクエリを融合するように設計されている。
   • 精緻化 (Refinement): ポストプロセスの精緻化メカニズムは、ユーザーのクリックを使用して、3D ガウスカーネルを介してマスク確率を局所的に調整する。

主な貢献

1. SAM2 の 3D 拡張: 本研究は、SAM2 フレームワークを 2D 自然画像から 3D 医学画像解析へと拡張することに成功し、スライスごとの近似を行うことなく、直接的なボリューム処理を可能にした。
2. ハイブリッドデコーダーアーキテクチャ: UNETR に着想を得たデコーダーの導入により、統一されたアーキテクチャ内で完全自動およびインタラクティブなセグメンテーションの両方を効果的にサポートし、局所的な微細詳細とグローバルなコンテキスト依存性の両方を捉えることが可能となった。
3. 包括的なデータセットのキュレーション: 著者らは、脳腫瘍、脳卒中、てんかんをカバーする大規模なデータセットをキュレーションし、幅広い臨床アプリケーションへの対応を可能にした。

結果
6 つの MRI モダリティをカバーする 7 つの公開データセット（BraTS-PED, BraTS-SSA, BraTS-MET, BraTS-MEN, ISLES2022, ATLAS v2.0, OPMRI）を用いて広範な実験が行われた。Brain-SAM は、タスク特化型モデル（nnU-Net, UNETR, Swin UNETR）および基盤モデル（SAM-Med3D, VISTA3D）と比較された。

• 全体的な性能: Brain-SAM は、すべてのベースラインと比較して優れた Dice スコアと低い 95% ハウスドルフ距離 (HD95) を達成した。具体的には、nnU-Net と比較して、てんかん病変で 22%、脳転移で 9%、髄膜腫で 6% の Dice 指標の向上を実現した。
• 小さな病変のセグメンテーション: 本モデルは、小さな病変（例：てんかん）のセグメンテーションにおいて顕著な優位性を示し、強力なベースラインモデルよりも 15%～18% 高い Dice スコアを達成した。また、他の手法と比較して、小さなターゲット（<5,000 ボクセル）に対する再現率と適合率のバランスをより良く維持した。
• アブレーション研究:
  • 3D-ViT バックボーンを 3D-Hiera に置き換えることで、小さな病変に対して Dice スコアが最大 14% 向上し、HD95 が大幅に減少した。
  • UNETR に着想を得たデコーダーを統合することで、未導入バージョンと比較して平均 Dice スコアが 8% 上昇した。
  • アダプターベースのファインチューニング戦略は、フルファインチューニングと同等の性能を達成しながら、非常にパラメータ効率が高いこと（フルファインチューニングの 72.02M に対し、調整可能なパラメータは 39.4M）が証明された。

意義と主張
論文は、Brain-SAM が下流の脳 MRI 解析タスクのための堅牢な事前学習済みモデルを提供すると主張している。不明瞭な境界や小さな病変サイズといった脳組織の特有の解剖学的および構造的な課題に対処することで、既存の汎用的な基盤モデルよりも信頼性が高く正確なセグメンテーションを実現している。著者らは、この能力が、特にてんかんや脳腫瘍のような疾患の早期診断や治療計画における、より情報に基づいた臨床的意思決定をサポートできると考えている。自動モードとプロンプト誘導モードの両方で機能する本モデルの能力は、高い精度を維持しながら手動のアノテーション作業を軽減することを目指している。著者らは、現在の研究は単一モダリティのデータに限定されていることを指摘しており、今後の研究では、病変の特性評価をさらに強化するためにマルチモダリティの統合を探求できる可能性があるとしている。