Improving Prostate Gland Segmentation Using Transformer based Architectures

論文概要

本論文は、T2 強調 MRI 画像を用いた前立腺の自動セグメンテーションにおいて、読者間の変動（Inter-reader variability）とドメインシフト（Cross-site-domain shift）という課題に対処するため、従来の CNN（3D U-Net）と比較してトランスフォーマーベースのモデル（UNETR と SwinUNETR）の有効性を検証した研究です。546 例の MRI ボリューム（2 人の専門読者による注釈付き）を用いた大規模なベンチマークにより、トランスフォーマーモデルが特に不均衡なデータや異なる注釈スタイルに対して高いロバスト性を示すことを実証しました。

1. 解決すべき課題 (Problem)

• 読者間の変動: 前立腺の輪郭描画は手作業であり、読者（放射線科医）によって描画のスタイルや基準にばらつきが生じます。これにより、ある読者のデータで学習したモデルが、別の読者のデータに対して性能が低下する問題が発生します。
• ドメインシフトと一般化: 単一の読者データセットや特定の施設で学習したモデルは、異なる読者や異なる条件下でのデータに対して一般化能力が不足しています。
• クラス不均衡と小規模な腺体: 前立腺は画像ボリュームの 3% 未満を占めることが多く、背景（クラス）に対して前景（前立腺）の比率が極端に低いクラス不均衡の問題があります。また、腺体のサイズが小さい場合、CNN は長距離の文脈（コンテキスト）を捉えるのに苦労します。
• 既存モデルの限界: 従来の 3D U-Net は局所的な特徴には優れていますが、長距離依存関係の捕捉が弱く、読者間の変動に対する適応性に限界がありました。

2. 手法と実験設計 (Methodology)

使用モデル

• 3D U-Net (Baseline): 従来の畳み込みニューラルネットワーク（CNN）ベースのモデル。
• UNETR: エンコーダーに純粋な Vision Transformer (ViT) を採用し、入力ボリュームをパッチに分割してグローバルな自己注意（Self-attention）を計算するモデル。
• SwinUNETR: エンコーダーに「シフトド・ウィンドウ（Shifted-window）」メカニズムを持つ Swin Transformer を採用。局所的なウィンドウ内での注意計算とウィンドウ間のシフトを組み合わせることで、グローバルな文脈を捉えつつ、メモリ使用量を線形に抑える効率的なモデル。

データセット

• 総数: 546 例の T2 強調 MRI ボリューム。
• 読者構成:
  • Reader#1 (R#1): 342 例（ProstateX チャンネルから取得）。
  • Reader#2 (R#2): 204 例（R#1 と同じ症例に対する別の専門家の注釈）。
  • 204 例は両読者によって注釈されており、読者間比較が可能です。

学習戦略 (3 つの実験)

1. 単一読者学習 (Single-reader training): 特定の読者のデータのみで学習し、他方の読者のデータでテスト。読者スタイルへの感度を評価。
2. 混合読者クロスバリデーション (Mixed-cohort cross-validation): R#1 と R#2 のデータを混合し、5 分割クロスバリデーションを実施。読者間一般化能力を評価。
3. 腺体サイズに基づくストレッステスト (Gland-size based stress test): 前立腺のサイズ（大・小）でデータを層化し、特定のサイズのデータ比率を変化させて学習。極端なクラス不均衡や形状のばらつきに対するロバスト性を評価。

最適化と評価

• ハイパーパラメータ: Optuna フレームワークを用いて、学習率、バッチサイズ、特徴量サイズを自動最適化。
• 損失関数: 重み付き Dice 損失とバイナリクロスエントロピー損失の組み合わせ（クラス不均衡対策）。
• 評価指標: Dice Similarity Coefficient (DSC)。統計的有意性の検定には Wilcoxon 符号付き順位和検定および Friedman 検定を使用。

3. 主要な結果 (Key Results)

全体的な性能

• SwinUNETR の優位性: 単一読者学習および混合学習の両方で、SwinUNETR が最も高い性能を示しました。
  • 単一読者学習（R#1 学習、R#1 テスト）: SwinUNETR (0.820) > 3D U-Net (0.825) ≈ UNETR (0.801)。
  • 混合学習（クロスバリデーション）: SwinUNETR は R#1 で 0.8583、R#2 で 0.8678 の平均 DSC を達成しました。
• UNETR の性能: 混合学習では 3D U-Net よりも優れ、SwinUNETR と同等の性能（R#1: 0.8517, R#2: 0.8640）を示しましたが、単一読者学習ではやや不安定でした。
• 3D U-Net: 混合学習でも性能向上が見られましたが、トランスフォーマーモデルに比べると、特に R#1 における性能が低く、読者間変動に対して脆弱でした。

腺体サイズによるロバスト性

• 大・小サイズへの適応: 腺体サイズが極端に偏ったデータセット（例：R#1 の小サイズデータを 10% まで削減）においても、SwinUNETR は高い安定性を維持しました。
  • 小サイズデータにおいて、SwinUNETR は DSC 0.85 前後を安定して維持しましたが、UNETR と 3D U-Net はデータ量の減少に伴い性能が顕著に低下しました（特に 3D U-Net は 0.78 まで低下）。
• 統計的有意性: Friedman 検定により、R#1 において SwinUNETR と他のモデル間の差は統計的に有意（p < 0.0001）であることが確認されました。

定性的評価

• SwinUNETR は、前立腺の輪郭が滑らかで、過剰セグメンテーションや過小セグメンテーションのアーティファクトが少なく、専門家の注釈と高い一致を示しました。
• UNETR は小サイズ腺体において過剰・過小セグメンテーションの傾向が見られ、クラス不均衡に敏感であることが示唆されました。

4. 主な貢献 (Key Contributions)

1. 包括的なベンチマーク: 同一の多読者 MRI データセット（546 例）を用いて、3D U-Net、UNETR、SwinUNETR を系統的に比較した初の研究です。
2. 読者間変動への耐性: トランスフォーマーベースのモデル、特に SwinUNETR が、異なる読者の注釈スタイルやドメインシフトに対して、CNN ベースのモデルよりも高いロバスト性と一般化能力を持つことを実証しました。
3. 不均衡データへの対応: 前立腺のサイズが小さい場合や、特定の読者データが極端に少ない状況でも、SwinUNETR がグローバル自己注意メカニズムにより安定した性能を維持できることを示しました。
4. 臨床応用への示唆: SwinUNETR は計算効率と精度のバランスが良く、GPU 上でリアルタイムに近い速度で動作するため、臨床現場での自動セグメンテーションの実用化に有望であることを示唆しています。

5. 意義と結論 (Significance & Conclusion)

本研究は、前立腺癌の早期発見と治療計画において重要な前立腺の自動セグメンテーション技術において、トランスフォーマーアーキテクチャ（特に SwinUNETR）が従来の CNN を凌駕することを示しました。

• 精度の向上: 3D U-Net ベースラインと比較して、DSC が最大で 5 ポイント向上しました。
• ノイズ耐性: グローバル自己注意メカニズムが、読者によるラベルノイズ（注釈のばらつき）やクラス不均衡の影響を軽減し、より頑健なモデルを実現しました。
• 将来展望: 本研究は、異なる施設や読者間で通用する「読者非依存（reader-agnostic）」なセグメンテーションモデルの構築に向けた重要な一歩であり、臨床ワークフローへの統合を加速させる可能性があります。

限定的な点として、本研究は T2 強調画像に限定されており、複数のスキャナやベンダーでの検証が今後の課題ですが、トランスフォーマーモデルの導入が医療画像解析の課題解決に有効であることが明確に示されました。