[论文解读] Estimating Uncertainty in Neural Networks for Cardiac MRI Segmentation: A Benchmark Study

Matthew Ng, Fumin Guo+5 more|arXiv (Cornell University)|Dec 31, 2020

Explainable Artificial Intelligence (XAI)被引用 7

一句话总结

本基准研究评估了在心脏MRI分割的U-Net模型中，不确定性估计方法——贝叶斯神经网络（BBB）、蒙特卡洛 dropout（MCD）、深度集成（Ensemble）和随机谱归一化（SSN）——的表现。结果表明，BBB 和 MCD 在噪声和失真条件下提供了可靠的不确定性量化，并保持了较高的分割精度，而深度集成在分布外场景下表现最为稳健。

ABSTRACT

Objective: Convolutional neural networks (CNNs) have demonstrated promise in automated cardiac magnetic resonance image segmentation. However, when using CNNs in a large real-world dataset, it is important to quantify segmentation uncertainty and identify segmentations which could be problematic. In this work, we performed a systematic study of Bayesian and non-Bayesian methods for estimating uncertainty in segmentation neural networks. Methods: We evaluated Bayes by Backprop, Monte Carlo Dropout, Deep Ensembles, and Stochastic Segmentation Networks in terms of segmentation accuracy, probability calibration, uncertainty on out-of-distribution images, and segmentation quality control. Results: We observed that Deep Ensembles outperformed the other methods except for images with heavy noise and blurring distortions. We showed that Bayes by Backprop is more robust to noise distortions while Stochastic Segmentation Networks are more resistant to blurring distortions. For segmentation quality control, we showed that segmentation uncertainty is correlated with segmentation accuracy for all the methods. With the incorporation of uncertainty estimates, we were able to reduce the percentage of poor segmentation to 5% by flagging 31--48% of the most uncertain segmentations for manual review, substantially lower than random review without using neural network uncertainty (reviewing 75--78% of all images). Conclusion: This work provides a comprehensive evaluation of uncertainty estimation methods and showed that Deep Ensembles outperformed other methods in most cases. Significance: Neural network uncertainty measures can help identify potentially inaccurate segmentations and alert users for manual review.

研究动机与目标

• 评估并比较多种不确定性估计技术在深度学习模型中用于心脏MRI分割的性能。
• 评估不同不确定性方法在各种图像失真（如噪声、模糊和几何变换）下维持分割精度的能力。
• 研究预测不确定性度量（如熵、互信息）与实际分割误差之间的相关性。
• 确定在临床影像场景中，哪种不确定性估计方法能提供最可靠且校准良好的不确定性估计。
• 基准比较每种不确定性方法在真实心脏MRI分割任务中的计算成本和推理效率。

提出的方法

• 在U-Net架构中训练了五种不确定性估计技术：贝叶斯神经网络（BBB）、蒙特卡洛 dropout（MCD），dropout率分别为0.1和0.5，深度集成（10个模型），以及随机谱归一化（SSN）。
• 训练过程中应用数据增强，包括随机旋转、平移和缩放，以提高模型鲁棒性。
• 使用Adam优化器并采用余弦退火学习率调度：初始值为1e-4，在第30个epoch后衰减至1e-5。
• 通过中心裁剪至160×160并使用数据集整体的均值和方差进行归一化，完成图像预处理。
• 在具有受控失真的测试图像上，使用预测熵、互信息以及不确定性感知指标（Dice_WS、ASSD_WS）测量不确定性。
• 在UK Biobank和ACDC数据集上评估模型性能，包括从UKBB到ACDC的零样本迁移。

实验结果

研究问题

• RQ1 在图像噪声逐渐增加的情况下，不同不确定性估计方法（BBB、MCD、Ensemble、SSN）在分割精度上的表现如何比较？
• RQ2 预测不确定性（如熵、互信息）与失真心脏MRI扫描中的实际分割误差之间的相关性有多大？
• RQ3 在从UK Biobank泛化到ACDC数据集时，哪种不确定性方法能提供最可靠且校准良好的不确定性估计？
• RQ4 各种方法的不确定性估计计算成本有何差异，特别是在推理时间和训练时长方面？
• RQ5 使用不确定性估计是否能提升模型对分布外失真（如模糊和几何形变）的鲁棒性？

主要发现

• 深度集成在干净的UK Biobank图像上实现了左心室（LV）分割最高的中位Dice分数（0.965，93.6–98.3）和最低的ASSD（0.74，0.43–1.12），优于普通U-Net和其他方法。
• BBB 和 MCD-0.1 在噪声和模糊条件下表现出更优的鲁棒性，相较于普通U-Net，保持了更高的像素级精度（如BBB：噪声水平1时为0.9929）和Dice分数（如BBB：噪声水平1时为0.947）。
• 互信息（MI）随失真水平显著增加——例如，SSN的MI从UKBB的0.32×10⁻³上升至UKBB→ACDC的6.89×10⁻³，表明不确定性与性能退化之间存在强烈相关性。
• 在ACDC数据集上，BBB 和 MCD-0.1 在互信息与ASSD之间表现出最高的斯皮尔曼等级相关性（r > 0.7），证明了其不确定性估计的可靠性与校准性。
• SSN 产生了最高的不确定性值（如UKBB→ACDC的MI为6.89×10⁻³），但分割精度最差（Dice LV：0.961），表明其在不确定性估计上存在过度自信。
• 统计检验确认了显著的性能差异：BBB 与普通U-Net（p < 0.0001）、MCD-0.1 与 MCD-0.5（p < 0.0001）、以及Ensemble 与 SSN（p < 0.0001）之间存在显著差异，凸显了各方法的可靠性差异。