[论文解读] 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

一句话总结

本文提出最大熵错误预测（MEEP）训练策略，通过仅对分类错误的像素施加低熵（过度自信）惩罚，提升医学图像分割模型的校准性能。通过在错误预测上最大化熵（采用熵或KL正则化），MEEP在不增加模型复杂度的前提下增强模型对不确定性的感知能力，显著改善了多种数据集和网络架构下的校准性能与分割性能。

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.

研究动机与目标

• 解决基于深度学习的医学图像分割中过度自信、校准不足的关键问题。
• 通过增强不确定性估计，特别是高不确定性、分类错误区域的估计，提升模型可靠性。
• 开发一种与网络架构无关的训练策略，不增加模型复杂度，且可与标准分割损失结合使用。
• 证明在评估临床人工智能模型时，校准指标与分割指标同等重要。
• 表明将正则化聚焦于错误预测，相比全局正则化，能获得更优的校准效果。

提出的方法

• 提出两种正则化项：一种直接在分类错误的像素上最大化熵，另一种使用KL散度项作为代理，以鼓励错误预测具有高不确定性。
• 仅在训练过程中对分类错误的像素应用正则化，确保聚焦于困难且不确定的样本。
• 将正则化整合到主损失函数中，不修改网络架构，也不增加推理成本。
• 使用基于熵的正则化，促使分类错误像素的预测概率分布趋向均匀（高熵），降低过度自信。
• 采用加权版本的预测掩码（$\hat{Y}_w$），将正则化聚焦于高不确定性且分类错误的像素。
• 可与标准分割损失（如Dice、交叉熵、焦点损失）结合使用，保持与现有训练流程的兼容性。

实验结果

研究问题

• RQ1 在分类错误的像素上惩罚低熵是否能改善医学图像分割中的模型校准？
• RQ2 仅对错误预测应用不确定性正则化，是否比全局正则化带来更好的校准效果？
• RQ3 MEEP是否可在不增加模型复杂度的前提下，应用于不同网络架构和分割损失函数？
• RQ4 通过MEEP提升的校准性能是否也带来更好的分割性能？
• RQ5 在可靠性与不确定性估计方面，MEEP与现有校准技术相比表现如何？

主要发现

• MEEP显著改善了模型校准性能，在WMH数据集上结合Dice或交叉熵损失时，Brier分数最高降低65%，ECE最高降低60%。
• 该方法实现了最先进的校准性能，在WMH数据集上使用KL正则化与交叉熵损失时，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等校准指标对于评估临床可靠性至关重要。