[論文レビュー] Generalised Wasserstein Dice Score for Imbalanced Multi-class Segmentation using Holistic Convolutional Networks

Lucas Fidon, Wen-Zhi Li+5 more|UCL Discovery (University College London)|Jul 3, 2017

Artificial Intelligence in Healthcare被引用数 16

ひとこと要約

本論文は、ラベル空間の木構造と包括的畳み込みニューラルネットワーク（HCNN）を用いてクラス間の意味的関係を組み込み、不均衡な多クラス脳腫瘍セグメンテーションを改善する、新たな一般化 Wasserstein Dice スコアを提案する。この手法は、特に希少な腫瘍亜型において、ベースラインの損失関数やアーキテクチャよりも高い Dice スコアとより意味的に整合性のあるセグメンテーションを達成する。

ABSTRACT

The Dice score is widely used for binary segmentation due to its robustness to class imbalance. Soft generalisations of the Dice score allow it to be used as a loss function for training convolutional neural networks (CNN). Although CNNs trained using mean-class Dice score achieve state-of-the-art results on multi-class segmentation, this loss function does neither take advantage of inter-class relationships nor multi-scale information. We argue that an improved loss function should balance misclassifications to favour predictions that are semantically meaningful. This paper investigates these issues in the context of multi-class brain tumour segmentation. Our contribution is threefold. 1) We propose a semantically-informed generalisation of the Dice score for multi-class segmentation based on the Wasserstein distance on the probabilistic label space. 2) We propose a holistic CNN that embeds spatial information at multiple scales with deep supervision. 3) We show that the joint use of holistic CNNs and generalised Wasserstein Dice scores achieves segmentations that are more semantically meaningful for brain tumour segmentation.

研究の動機と目的

• 腫瘍亜型間の意味的関係を活用することで、多クラス脳腫瘍セグメンテーションにおけるクラス不均衡を解消すること。
• 新しい包括的 CNN アーキテクチャを用いて多スケール空間特徴を統合することで、セグメンテーションのロバスト性を向上させること。
• 階層的ラベル構造を符号化する微分可能な損失関数を開発し、より良い一般化性能を達成すること。
• 意味的感度を持つ損失関数と多スケール監視を組み合わせることで、優れたセグメンテーション性能が得られることを実証すること。

提案手法

• ラベル空間における事前定義された距離行列を用いて意味的関係を符号化する、ラベル確率ベクトル間の Wasserstein 距離に基づく一般化 Dice スコアを提案する。
• 複数スケールの予測を生成し、各スケールでディープ監視を適用することで特徴学習を向上させる、包括的畳み込みニューラルネットワーク（HCNN）を設計する。
• 訓練の安定性を高めるために、2段階の訓練戦略を採用：まず平均ソフト Dice 損失で事前学習し、その後、提案された Wasserstein Dice 損失で微調整する。
• 腫瘍タイプ間の解剖学的・病理学的関係を反映した手動で割り当てられたエッジ重みを有する木構造ラベル空間を用いる（例：壊死巣、強化しない、強化する）。
• 全クラスの確率的ラベルマップを同時に比較するために Wasserstein 距離を適用し、クラス間の誤分類ペナルティを統合的に最適化可能にする。
• 学習可能または固定の統合戦略を用いて多スケール予測を統合し、高解像度の詳細と広い受容 field のコンテキストを組み合わせる。

実験結果

リサーチクエスチョン

• RQ1 Wasserstein ベースの Dice スコアの一般化は、不均衡な多クラス医療画像セグメンテーションにおいて性能を向上させるか？
• RQ2 距離行列を用いた階層的ラベル関係の組み込みが、モデルの一般化性能と意味的整合性に与える影響は何か？
• RQ3 ディープ監視を伴う多スケール特徴学習は、脳腫瘍 MRI におけるセグメンテーション精度をどの程度向上させるか？
• RQ4 2段階の訓練戦略（平均 Dice での事前学習 followed by Wasserstein Dice での微調整）は、単独で使用するいずれの損失関数でも、収束性と性能が優れているか？

主な発見

• 提案された一般化 Wasserstein Dice 損失 ($\mathcal{L}_{D^{M_{\text{tree}}}}$) は、壊死巣や強化しない腫瘍などの希少クラスにおいて、平均ソフト Dice 損失よりも高い Dice スコアを達成した。
• 多スケール監視とディープ監視を備えた包括的 CNN（HCNN）は、標準的な U-Net 形式アーキテクチャよりも一貫性があり、より正確なセグメンテーションを生成した。
• HCNN の予測統合により、全腫瘍 Dice スコアが 89% (±6)、コア腫瘍が 73% (±22)、強化腫瘍が 74% (±23)、$D^{M_{\text{tree}}}$ が 87% (±4) を達成し、個々のスケールやベースラインを上回った。
• 平均 Dice 損失での事前学習 followed by Wasserstein Dice 損失での微調整 ($\mathcal{L}_{D^{M_{\text{tree}}}-\text{PT}}$) により、訓練が安定化し、直接 $\mathcal{L}_{D^{M_{\text{tree}}}}$ で学習する場合に比べ、真正陽性率が向上した。
• スケール 1 は受容 field が不十分なため低い Dice スコアを示したが、スケール 3 と 4 は解像度が低下したため、細い腫瘍領域を正しく捉えられず、スケールと解像度のトレードオフが顕在化した。
• 定性的な結果から、強化腫瘍領域（最小で最も重要な構造）を正しくセグメンテーションできたのは、統合予測とスケール 2 のみであった。