Confidence Calibration and Predictive Uncertainty Estimation for Deep Medical Image Segmentation

Abstract

Fully convolutional neural networks (FCNs), and in particular U-Nets, have achieved state-of-the-art results in semantic segmentation for numerous medical imaging applications. Moreover, batch normalization and Dice loss have been used successfully to stabilize and accelerate training. However, these networks are poorly calibrated i.e. they tend to produce overconfident predictions both in correct and erroneous classifications, making them unreliable and hard to interpret. In this paper, we study predictive uncertainty estimation in FCNs for medical image segmentation. We make the following contributions: 1) We systematically compare cross entropy loss with Dice loss in terms of segmentation quality and uncertainty estimation of FCNs; 2) We propose model ensembling for confidence calibration of the FCNs trained with batch normalization and Dice loss; 3) We assess the ability of calibrated FCNs to predict segmentation quality of structures and detect out-of-distribution test examples. We conduct extensive experiments across three medical image segmentation applications of the brain, the heart, and the prostate to evaluate our contributions. The results of this study offer considerable insight into the predictive uncertainty estimation and out-of-distribution detection in medical image segmentation and provide practical recipes for confidence calibration. Moreover, we consistently demonstrate that model ensembling improves confidence calibration.

A Quantitative Comparison of Epistemic Uncertainty Maps Applied to Multi-Class Segmentation

“…The authors recommend the use of this uncertainty measure compared to variance based ones since it is more interpretable as the range of uncertainty values is between 0 and 1 (0 very certain, 1 very uncertain). In another work, Mehrtash et al (2019) compared calibrated and uncalibrated segmenta-…”

Local Temperature Scaling for Probability Calibration

“…Other Approaches. Mehrtash et al [38] found that model ensembling improves confidence calibration for medical image segmentation. Karimi et al [26] showed that multi-task learning could yield better-calibrated predictions than dedicated models trained separately.…”

“…The following key limitations of existing calibration approaches still need to be addressed: (1) Most of the probability calibration approaches are designed for classifications, thus are not guaranteed to work well for semantic segmentation; (2) Most methods are designed to work for binary classifications and approach multi-class problems by a decomposition into k one-vs-rest binary calibrations (where k denotes the number of classes). However, such a decomposition does not guarantee overall calibration (only for the individual subproblems before normalization) and the classification accuracy of the trained model may change after calibration as the probability order of labels may change; (3) Limited work discusses the probability calibration of semantic segmentation, but this work either only applies to specific types of models (e.g., Bayesian neural networks [24]) or only implicitly improves calibration performance (e.g., via model ensembling [38] or multi-task learning [26]).…”

Uncertainty Quantification and Resource-Demanding Computer Vision Applications of Deep Learning

“…In the past years, UQ for image segmentation and object detection has been initiated in a series of works [16][17][18][19][20][21] focusing on false positive instances and false negative instances 22,23 , respectively, we also refer to a survey 17 . While many works focus on street scene recognition for autonomous driving, 16,24 cover medical image segmentation applications as well.…”