[論文レビュー] A Mixed Focal Loss Function for Handling Class Imbalanced Medical Image Segmentation.

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

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

ひとこと要約

本稿では、修正されたFocal損失とFocal Dice損失を組み合わせたMixed Focal損失関数を提案し、腫瘍が小さい小規模な背景領域と比較してクラス不均衡が著しい医療画像セグメンテーションの性能を向上させる。KiTS19およびBraTS20データセットにおける二値および多クラスセグメンテーションにおいて、標準損失関数および構成要素損失関数を上回る優れた再現率-適合率のバランスと頑健性を達成した。

ABSTRACT

Automatic segmentation methods are an important advancement in medical imaging analysis. Machine learning techniques, and deep neural networks in particular, are the state-of-the-art for most automated medical image segmentation tasks, ranging from the subcellular to the level of organ systems. Issues with class imbalance pose a significant challenge irrespective of scale, with organs, and especially with tumours, often occupying a considerably smaller volume relative to the background. Loss functions used in the training of segmentation algorithms differ in their robustness to class imbalance, with cross entropy-based losses being more affected than Dice-based losses. In this work, we first experiment with seven different Dice-based and cross entropy-based loss functions on the publicly available Kidney Tumour Segmentation 2019 (KiTS19) Computed Tomography dataset, and then further evaluate the top three performing loss functions on the Brain Tumour Segmentation 2020 (BraTS20) Magnetic Resonance Imaging dataset. Motivated by the results of our study, we propose a Mixed Focal loss function, a new compound loss function derived from modified variants of the Focal loss and Focal Dice loss functions. We demonstrate that our proposed loss function is associated with a better recall-precision balance, significantly outperforming the other loss functions in both binary and multi-class image segmentation. Importantly, the proposed Mixed Focal loss function is robust to significant class imbalance. Furthermore, we showed the benefit of using compound losses over their component losses, and the improvement provided by the focal variants over other variants.

研究の動機と目的

• 腫瘍が小さい背景領域と比較してクラス不均衡が著しい医療画像セグメンテーションにおける課題に対処すること。
• 実世界の医療画像データセットにおけるDiceベースおよび交差エントロピーに基づく損失関数の頑健性と性能を評価すること。
• 顕著なクラス不均衡下での既存損失関数の限界を特定し、より効果的な複合損失の開発を促すこと。
• Focal損失とFocal Dice損失の長所を組み合わせた、新たなMixed Focal損失関数を提案し、セグメンテーション性能を向上させること。

提案手法

• 著者らはまず、KiTS19 CTデータセットを用いて、交差エントロピーおよびDiceベースの変種を含む7つの既存損失関数を評価し、上位の性能を示す候補を同定した。
• 次に、上位3つの損失関数を選択し、BraTS20 MR画像データセットを用いてモodalの観点から一般化性能を検証した。
• 提案されたMixed Focal損失は、修正されたFocal損失と修正されたFocal Dice損失を組み合わせて構築され、容易なネガティブサンプルの重みを低減し、誤分類の難しい例に学習を集中させるように設計された。
• 複合損失関数は、モデル学習中にエンドツーエンドで最適化され、二値および多クラスセグメンテーションタスクにおける再現率と適合率のバランスを最適化するようにハイパーパrameterを調整した。
• この手法は、容易で背景優勢なサンプルの影響を低減するためのフォーカルモodulationを活用し、腫瘍のような希少で重要な構造の学習を向上させた。

実験結果

リサーチクエスチョン

• RQ1 医療画像セグメンテーションにおいて顕著なクラス不均衡が生じる状況下で、標準的な交差エントロピーおよびDiceベースの損失関数はどの程度の性能を示すか？
• RQ2 Focal損失とFocal Dice損失を組み合わせた複合損失関数は、個々の構成要素よりも優れた再現率-適合率のバランスを達成できるか？
• RQ3 提案されたMixed Focal損失は、CTおよびMRのような異なる画像モodal間で一般化可能か？
• RQ4 標準損失関数と比較して、フォーカルメカニズムは腫瘍のような小規模で希少な構造の認識性能をどの程度向上させるか？
• RQ5 クラス不均衡が著しいセグメンテーションタスクにおいて、複合損失関数はその構成要素損失関数よりも頑健で効果的か？

主な発見

• Mixed Focal損失関数は、KiTS19およびBraTS20データセットにおける二値および多クラスセグメンテーションタスクにおいて、すべてのベースライン損失関数を顕著に上回った。
• 提案された損失関数は再現率と適合率のバランスを著しく改善し、特に標準損失関数が見逃しがちな小規模で希少な腫瘍領域の検出性能が向上した。
• 複合損失関数は顕著なクラス不均衡に対して頑健であり、腫瘍領域が全体積の1％未満である場合でも高い性能を維持した。
• 損失関数のフォーカル変種は、非フォーカル対応のものよりも一貫して優れた性能を示し、誤分類の難しいサンプルに注目することで利益が得られることを確認した。
• 本研究では、複合損失を用いることで個々の構成損失関数を用いる場合に比べて測定可能な性能向上が得られることを確認した。これにより、Mixed Focalアプローチの設計が妥当であることが裏付けられた。