[論文レビュー] Combo Loss: Handling Input and Output Imbalance in Multi-Organ Segmentation

Saeid Asgari Taghanaki, Yefeng Zheng+6 more|arXiv (Cornell University)|May 7, 2018

Radiomics and Machine Learning in Medical Imaging被引用数 4

ひとこと要約

本論文では、カリキュラム学習フレームワーク内でDice損失と重み付き交差エントロピーを組み合わせることで、入力と出力の両方のクラス不均衡を同時に解消する、新しい損失関数であるCombo Lossを提案する。この手法により、トレーニング中に偽陽性/偽陰性のトレードオフを明示的に制御可能となり、PET、MRI、超音波画像データセットにおいて収束が速く、最先端の性能を達成する。Focal Lossを上回る最高で10.8%のDiceスコア向上を実現した。

ABSTRACT

Simultaneous segmentation of multiple organs from different medical imaging modalities is a crucial task as it can be utilized for computer-aided diagnosis, computer-assisted surgery, and therapy planning. Thanks to the recent advances in deep learning, several deep neural networks for medical image segmentation have been introduced successfully for this purpose. In this paper, we focus on learning a deep multi-organ segmentation network that labels voxels. In particular, we examine the critical choice of a loss function in order to handle the notorious imbalance problem that plagues both the input and output of a learning model. The input imbalance refers to the class-imbalance in the input training samples (i.e., small foreground objects embedded in an abundance of background voxels, as well as organs of varying sizes). The output imbalance refers to the imbalance between the false positives and false negatives of the inference model. In order to tackle both types of imbalance during training and inference, we introduce a new curriculum learning based loss function. Specifically, we leverage Dice similarity coefficient to deter model parameters from being held at bad local minima and at the same time gradually learn better model parameters by penalizing for false positives/negatives using a cross entropy term. We evaluated the proposed loss function on three datasets: whole body positron emission tomography (PET) scans with 5 target organs, magnetic resonance imaging (MRI) prostate scans, and ultrasound echocardigraphy images with a single target organ i.e., left ventricular. We show that a simple network architecture with the proposed integrative loss function can outperform state-of-the-art methods and results of the competing methods can be improved when our proposed loss is used.

研究の動機と目的

• 深層多臓器セグメンテーションにおける、入力不均衡（少数の前景ボクセル）と出力不均衡（偽陽性対偽陰性）の二重の課題に対処すること。
• トレーニング中に偽陽性と偽陰性のトレードオフを明示的かつ学習可能な形で制御できる損失関数の開発。
• PET、MRI、超音波画像を含む多様なモダリティにおける、極めて不均衡な医療画像データに対するセグメンテーション性能の向上。
• アーキテクチャの変更なしに既存のアーキテクチャを強化できる、即挿し可能な損失関数の提供により、一般化性能と収束性の向上。
• 後処理への依存を減らすために、微分可能で適応的な損失を通じて、偽陽性・偽陰性を直接ペナルティ化すること。

提案手法

• Combo Lossは、トレーニングの安定化と悪くない局所最適解の回避を目的として、Dice類似係数（DSC）と重み付き交差エントロピー（WCE）を1つの微分可能な損失関数に統合する。
• 損失関数は以下の式で定式化される：L_combo = α × (1 - DSC) + (1 - α) × WCE、ここでαはDSC項とWCE項のバランスを制御する。
• WCE項に第2のハイパーパrameter βを導入し、偽陽性と偽陰性ペナルティのトレードオフを明示的に制御する。
• 最適化の安定性と収束速度の向上を図るため、WCE項の影響を段階的に増加させるカリキュラム学習を採用する。
• アーキテクチャの変更なしに、エンコーダ・デコーダネットワーク（例：3D U-Net、V-Net）にエンドツーエンドで適用可能。
• ハイパーパrameter α と β は、各データセットでの性能最適化を目的に交差検証により調整される。

実験結果

リサーチクエスチョン

• RQ1 統合された損失関数は、多臓器セグメンテーションにおける入力と出力の両方の不均衡を効果的に処理できるか？
• RQ2 Dice損失と重み付き交差エントロピーを組み合わせることで、不均衡な医療画像データに対するモデルの収束性と性能がどのように向上するか？
• RQ3 損失関数内の学習可能なハイパーパrameterを通じて、偽陽性と偽陰性のトレードオフをどの程度制御可能か？
• RQ4 提案されたCombo Lossは、交差エントロピー、Dice損失、Focal Lossといった既存の損失関数を上回る性能を、多様な画像モダリティで示せるか？
• RQ5 Combo Lossは、アーキテクチャの変更なしに、既存のセグメンテーションネットワークの即挿し可能な代替損失関数として利用可能か？

主な発見

• PETデータセットでは、標準的な交差エントロピー損失と比較して、偽陽性率が34.6%低下し、偽陰性率が27.8%低下した。
• MRIデータセットでは、3D V-Netを用いた場合、Focal Lossと比較してDiceスコアが2.3%向上し、偽陰性率が47.4%低下した。
• 超音波データセットでは、3D U-Netに適用した場合、Focal Lossを上回る最高で10.8%のDiceスコア向上を達成した。
• 偽陽性と偽陰性を最小化する最適なβ値は、超音波データで0.7、MRIデータで0.6であった。これはハイパーパrameter制御の有効性を示している。
• U-NetとV-Netの元の損失関数を置き換えた場合、MRIではDiceスコアが0.87から0.91に、超音波では0.85から0.92に上昇し、一貫した向上が確認された。
• 標準的な交差エントロピー損失よりも収束が速く、ハウスドルフ距離における外れ値性能も低減した。これは、より頑健なセグメンテーションを実現していることを示している。