[論文レビュー] Maximum Entropy on Erroneous Predictions (MEEP): Improving model calibration for medical image segmentation

Agostina J. Larrazabal, César E. Martínez+2 more|arXiv (Cornell University)|Dec 22, 2021

Explainable Artificial Intelligence (XAI)被引用数 5

ひとこと要約

本稿では、誤分類された画素での低エントロピー（過信）にのみペナルティを課すことで、医療画像セグメンテーションにおけるモデルのキャリブレーションを向上させる、新しいトレーニング戦略であるMaximum Entropy on Erroneous Predictions (MEEP) を提案する。誤った予測におけるエントロピーを最大化するか、KL正則化を用いることで、モデルの複雑さを増すことなく不確実性の認識を高め、複数のデータセットおよびアーキテクチャにおいて顕著に改善されたキャリブレーションとセグメンテーション性能を達成する。

ABSTRACT

Modern deep neural networks achieved remarkable progress in medical image segmentation tasks. However, it has recently been observed that they tend to produce overconfident estimates, even in situations of high uncertainty, leading to poorly calibrated and unreliable models. In this work we introduce Maximum Entropy on Erroneous Predictions (MEEP), a training strategy for segmentation networks which selectively penalizes overconfident predictions, focusing only on misclassified pixels. Our method is agnostic to the neural architecture, does not increase model complexity and can be coupled with multiple segmentation loss functions. We benchmark the proposed strategy in two challenging segmentation tasks: white matter hyperintensity lesions in magnetic resonance images (MRI) of the brain, and atrial segmentation in cardiac MRI. The experimental results demonstrate that coupling MEEP with standard segmentation losses leads to improvements not only in terms of model calibration, but also in segmentation quality.

研究の動機と目的

• ディープラーニングベースの医療画像セグメンテーションにおける過信、かつキャリブレーションが不十分な予測という深刻な問題に対処すること。
• 特に高不確実性で誤分類された領域において不確実性推定を向上させることで、モデルの信頼性を高めること。
• アーキテクチャに依存しない戦略を構築し、モデルの複雑さを増すことなく、標準的なセグメンテーション損失と組み合わせられること。
• 臨床的AIモデルの評価において、セグメンテーション指標に加えてキャリブレーション指標が不可欠であることを示すこと。
• 正則化を誤った予測に限定することで、グローバルに正則化を適用するのと比較して、より良いキャリブレーションが得られることを示すこと。

提案手法

• 誤分類された画素におけるエントロピーを直接最大化する正則化項と、誤った予測における高い不確実性を促進するためのKLダイバージェンス項の2つの正則化項を提案する。
• トレーニング中に誤分類された画素にのみ正則化を適用することで、困難で不確実なケースに焦点を当てる。
• ネットワークアーキテクチャを変更せず、推論コストを増加させない形で、正則化をメイン損失関数に統合する。
• エントロピーに基づく正則化を用い、誤分類された画素の予測確率分布を一様（高エントロピー）に近づけることで、過信を低減する。
• 予測マスクの重み付きバージョン（$\hat{Y}_w$）を用い、不確実性が高く誤分類された画素に正則化を集中させる。
• Dice、クロスエントロピー、フォーカル損失などの標準的なセグメンテーション損失と組み合わせられ、既存のトレーニングパイプラインと互換性を持つ。

実験結果

リサーチクエスチョン

• RQ1 誤分類された画素における低エントロピーにペナルティを課すことで、医療画像セグメンテーションにおけるモデルのキャリブレーションが向上するか？
• RQ2 誤った予測にのみ不確実性正則化を適用することで、グローバル正則化と比較して、より良いキャリブレーションが得られるか？
• RQ3 MEEPはモデルの複雑さを増すことなく、さまざまなアーキテクチャやセグメンテーション損失関数に適用可能か？
• RQ4 MEEPによるキャリブレーションの向上は、セグメンテーション性能の向上にもつながるか？
• RQ5 MEEPは、既存のキャリブレーション技術と比較して、信頼性および不確実性推定の観点で優れているか？

主な発見

• MEEPはモデルのキャリブレーションを顕著に向上させ、Diceまたはクロスエントロピー損失を組みわせた場合、WMHデータセットにおいてBrierスコアを最大65％、ECEを最大60％低減した。
• 本手法は、状態を凌駕するキャリブレーション性能を達成し、KL正則化とクロスエントロピー損失を組みわせた場合、WMHデータセットでBrierスコア（0.434 × 10⁻³）とECE（0.174）で最高の結果を記録した。
• クロスエントロピー損失を組みわせたMEEPを用いることで、WMHではDice係数が最大0.016（1.6％）向上、LAでは0.013（1.4％）向上した。
• 提案手法はバックボーンに依存しない：UNetおよびResUNetの両アーキテクチャで同様の向上が観察され、汎用性が確認された。
• 定性的な結果から、MEEPでトレーニングされたモデルは、難易度の高いセグメンテーション領域において、0.5に近い確率マップを生成し、不確実性が高まっていることが示された。
• 本研究では、Diceのようなセグメンテーション指標にのみ依存することは不十分であることが示された。臨床的信頼性を評価するには、BrierスコアやECEのようなキャリブレーション指標が不可欠である。