[論文レビュー] Shape-aware Meta-learning for Generalizing Prostate MRI Segmentation to Unseen Domains

Quande Liu, Qi Dou+1 more|arXiv (Cornell University)|Jul 4, 2020

Domain Adaptation and Few-Shot Learning参考文献 30被引用数 12

ひとこと要約

本論文は、メタ最適化中に形状のcompactさと境界の滑らかさを明示的にモデル化することで、前立腺MRI画像のセグメンテーションにおけるドメイン一般化を向上させる、新しいメタラーニングフレームワークであるShape-aware Meta-learning (SAML)を提案する。エピソード学習を用いてドメインシフトをシミュレートし、2つの形状に配慮した損失関数を導入することで、6つの未学習の臨床施設においてDiceスコアおよびASD指標で一貫した向上を達成し、最先端の性能を実現した。

ABSTRACT

Model generalization capacity at domain shift (e.g., various imaging protocols and scanners) is crucial for deep learning methods in real-world clinical deployment. This paper tackles the challenging problem of domain generalization, i.e., learning a model from multi-domain source data such that it can directly generalize to an unseen target domain. We present a novel shape-aware meta-learning scheme to improve the model generalization in prostate MRI segmentation. Our learning scheme roots in the gradient-based meta-learning, by explicitly simulating domain shift with virtual meta-train and meta-test during training. Importantly, considering the deficiencies encountered when applying a segmentation model to unseen domains (i.e., incomplete shape and ambiguous boundary of the prediction masks), we further introduce two complementary loss objectives to enhance the meta-optimization, by particularly encouraging the shape compactness and shape smoothness of the segmentations under simulated domain shift. We evaluate our method on prostate MRI data from six different institutions with distribution shifts acquired from public datasets. Experimental results show that our approach outperforms many state-of-the-art generalization methods consistently across all six settings of unseen domains.

研究の動機と目的

• 一つのスキャナーやプロトコルで学習したモデルが未学習ドメインで失敗するという、前立腺MRI画像セグメンテーションにおけるドメインシフトの課題に対処すること。
• 訓練時にターゲットドメインのデータを必要とせず、新規の臨床環境に直接展開可能なモデルの一般化を向上させること。
• 未学習ドメインにおけるセグメンテーション出力の欠陥、例えば不完全な形状や曇った境界の問題を克服すること。
• 形状固有のインダクティブバイアス（compactさと滑らかさ）をメタラーニングに統合し、医療画像セグメンテーションにおけるより良い一般化を実現すること。
• モデルに依存しない勾配ベースのメタラーニングアプローチを構築し、複数施設のMRIデータにおける分布シフトに対してより頑健な性能を発揮させること。

提案手法

• 勾配ベースのメタラーニングフレームワークを採用し、エピソード学習を用いて、ソースドメインをメタトレインおよびメタテストセットに分割することでドメインシフトをシミュレートする。
• 形状のcompactさを促進する損失関数（$\mathcal{L}_{compact}$）を導入し、ドメインシフト下でも完全で閉じた形状を維持するようにセグメンテーションマスクを制御する。
• 潜在空間におけるドメイン不変な境界埋め込みを促進する形状滑らかさ損失（$\mathcal{L}_{smooth}$）を設計し、境界の明確さを向上させる。
• 標準のセグメンテーション損失に加え、$\mathcal{L}_{compact}$ と $\mathcal{L}_{smooth}$ を組み合わせたメタ最適化目的関数を採用し、多様な画像プロトコルにわたる頑健な特徴を学習可能にする。
• モデルに依存しないアプローチであり、任意のセグメンテーションバックボーンに適用可能であるため、柔軟かつ広範に展開可能である。
• トレーニングは、メタトレインおよびメタテストタスクの間で交互に最適化を実行し、メタアップデートの過程でシミュレートされたドメインシフトにモデルが適応するようにする。

実験結果

リサーチクエスチョン

• RQ1 明示的な形状に配慮した正則化を施したメタラーニングが、未学習の臨床ドメインにおける前立腺MRI画像セグメンテーションの一般化を改善できるか？
• RQ2 形状のcompactさと境界の滑らかさの制約が、ドメインシフト下でのセグメンテーション性能にどのように影響するか？
• RQ3 メタ最適化に形状のインダクティブバイアスを統合することで、標準のメタラーニングおよびデータ集約ベースラインを上回る性能が得られるか？
• RQ4 ソースドメインの数が、既存の手法と比較して、提案手法の一般化性能にどのように影響するか？
• RQ5 形状に配慮した損失関数が、未学習ドメインにおける不完全なセグメンテーションやぼやけた境界といった一般的な失敗モードを緩和できるか？

主な発見

• SAMLは、全6施設の平均Diceスコアが87.67に達し、ベースラインのDeepAll（85.52）および他の最先端手法を顕著に上回った。
• 最も挑戦的とされる施設（施設F）では、SAMLがDiceスコア88.34を達成し、次に良い性能を示した手法（MASF、86.55）を1.79ポイント上回った。
• 両方の損失関数$\mathcal{L}_{compact}$ と $\mathcal{L}_{smooth}$ を組み込むことで、全6施設で一貫した改善が得られ、ベースライン比でASDが1.22ポイント低下した。
• アブレーションスタディの結果、形状に配慮した両損失関数が独立してかつ相乗的に寄与しており、特に$\mathcal{L}_{smooth}$ が最大の性能向上をもたらした。
• ソースドメインの数が増加するにつれて、SAMLはDeepAllを常に上回り、安定的かつスケーラブルな一般化性能を維持した。
• 施設内での訓練では、6施設中の4施設でSAMLの性能が劣っており、複数ドメインのデータが限られた局所データよりも、頑健なモデル学習に効果的であることが示された。