[論文レビュー] Unified Focal loss: Generalising Dice and cross entropy-based losses to handle class imbalanced medical image segmentation

Michael Yeung, Evis Sala+2 more|arXiv (Cornell University)|Feb 8, 2021

Retinal Imaging and Analysis被引用数 9

ひとこと要約

本稿では、医療画像セグメンテーションにおけるクラス不均衡をより効果的に扱うために、Dice損失と交差エントロピー損失を一般化する階層的フレームワークであるUnified Focal Lossを提案する。Focal Lossの原則を異なる可微分で適応的な重み付けと組み合わせることで、5つの多様で不均衡なデータセットにおいて一貫した性能向上を達成し、二値および多クラスタスクの両方で標準的および複合損失関数を上回る性能を発揮する。

ABSTRACT

Automatic segmentation methods are an important advancement in medical image analysis. Machine learning techniques, and deep neural networks in particular, are the state-of-the-art for most medical image segmentation tasks. Issues with class imbalance pose a significant challenge in medical datasets, with lesions often occupying a considerably smaller volume relative to the background. Loss functions used in the training of deep learning algorithms differ in their robustness to class imbalance, with direct consequences for model convergence. The most commonly used loss functions for segmentation are based on either the cross entropy loss, Dice loss or a combination of the two. We propose the Unified Focal loss, a new hierarchical framework that generalises Dice and cross entropy-based losses for handling class imbalance. We evaluate our proposed loss function on five publicly available, class imbalanced medical imaging datasets: CVC-ClinicDB, Digital Retinal Images for Vessel Extraction (DRIVE), Breast Ultrasound 2017 (BUS2017), Brain Tumour Segmentation 2020 (BraTS20) and Kidney Tumour Segmentation 2019 (KiTS19). We compare our loss function performance against six Dice or cross entropy-based loss functions, across 2D binary, 3D binary and 3D multiclass segmentation tasks, demonstrating that our proposed loss function is robust to class imbalance and consistently outperforms the other loss functions. Source code is available at: https://github.com/mlyg/unified-focal-loss

研究の動機と目的

• 画像ボリュームのわずかな割合を占める病変や臓器に起因する、医療画像セグメンテーションにおける継続的なクラス不均衡の課題に対処すること。
• 標準的交差エントロピー損失およびDice損失の限界を克服すること。これらはクラス不均衡に対して感受性が高く、収束不良やバイアスのかかった予測を引き起こす可能性がある。
• Diceおよび交差エントロピーに基づく損失を一般化する統一的かつ可微分なフレームワークを構築し、不均衡下でもより頑健で高い性能を実現すること。
• 2D二値、3D二値、3D多クラスセグメンテーションを含む、多様な医療画像処理タスクにおいて、提案された損失関数の一貫した優位性を実証すること。
• 既存のU-Netおよび類似アーキテクチャに構造的変更を加えずに容易に統合可能な、柔軟で即席の損失関数を提供すること。

提案手法

• 異なる可微分で適応的な重み付けメカニズムを介して、Dice損失と交差エントロピー損失を統合する階層的損失フレームワークを提案する。
• クラス固有の動的重み付けを適用し、容易なネガティブ例を軽減し、ハード例に焦点を当てる一般化されたFocal損失部を導入する。
• 修正されたDice損失とFocal交差エントロピー損失の重み付き和としてUnified Focal Lossを定式化し、学習可能なまたは固定のハイパーパrameterを用いる。
• 集中効果を制御するためのガンマ（γ）パラメータを用い、多クラスタスクでは単一のγ=0.5を全クラスに適用して簡素化する。
• U-Netおよび類似のエンコーダ・デコーダアーキテクチャにおける標準的なバックプロパゲーションと互換性があり、微分可能であることを保証する。
• 標準的なU-Netおよび3D U-Netアーキテクチャを用いて、5つの公開データセットにおける2Dおよび3Dセグメンテーションネットワークの学習にこの損失関数を適用する。

実験結果

リサーチクエスチョン

• RQ1 クラス不均衡な医療画像セグメンテーションにおいて、Dice損失と交差エントロピー損失を一般化しつつ、頑健性を保つ統一された損失関数を設計できるか？
• RQ2 多様で不均衡な医療画像データセットにおいて、Unified Focal Lossは標準的Dice損失、交差エントロピー損失、および複合損失（例：Combo損失）と比較して、性能に優れているか？
• RQ3 腫瘍が大きな臓器に比べて極めて小さいような高不均衡状況、例えば小規模な腫瘍のセグメンテーションにおいて、提案された損失関数は安定した勾配を維持し、収束を改善するか？
• RQ4 複数の希少クラスを含む多クラスセグメンテーションタスク、例えば腎腫瘍セグメンテーションにおいて、Unified Focal Lossは一般化可能か？
• RQ5 異なる画像モodalitiesおよびネットワークアーキテクチャにおいて、Unified Focal Lossの性能向上は一貫しているか？

主な発見

• Unified Focal Lossは、CVC-ClinicDB、DRIVE、BUS2017、BraTS20、KiTS19の5つのすべてのデータセットにおいて、Dice、交差エントロピー、Combo損失、Tversky損失を含む6つのベースライン損失関数を一貫して上回った。
• 極めて不均衡なKiTS19データセットでは、腎臓および腎腫瘍セグメンテーションの両方で最高のDSC（Dice Similarity Coefficient）を達成し、それどころかCombo損失でさえも上回った。
• BraTS20データセットでは、Dice損失およびTversky損失が不安定な勾配のため失敗したが、Unified Focal Lossは安定した学習を維持し、優れたセグメンテーション品質を達成した。
• 定性的な結果から、Unified Focal Lossを用いて生成されたセグメンテーションは、他の損失関数と比較して、より鋭い境界線と、より少ない偽陽性・偽陰性を示した。
• 他の損失関数と比較して、学習時間に顕著な差が認められず、計算効率が高く、実用的かつ導入可能であることを示した。
• 単一のγ値を全クラスに適用したにもかかわらず、他の複合損失関数よりも優れた性能を達成したため、ハイパーパrameter設定に対して頑健であることが示唆された。