Confidence Calibration and Predictive Uncertainty Estimation for Deep Medical Image Segmentation

Alireza Mehrtash; William M. Wells; Clare M. Tempany; Purang Abolmaesumi; Tina Kapur

Source-linked AI summary

Confidence Calibration and Predictive Uncertainty Estimation for Deep Medical Image Segmentation

Alireza Mehrtash, William M. Wells, Clare M. Tempany, Purang Abolmaesumi, Tina Kapur

arXiv:1911.13273v2Nov 29, 2019eess.IVcs.CV

TL;DR

Medical FCNs often produce overconfident probabilities, motivating reliable uncertainty estimation and calibration. The paper compares loss functions, uses model ensembling, and evaluates segment-level entropy across three MRI tasks; ensembling consistently improves calibration, while entropy tracks segmentation quality and shows potential for out-of-distribution detection.

Problem
Medical-image FCNs can be poorly calibrated and overconfident, while predictive uncertainty and out-of-distribution detection remain insufficiently studied.
Method
The study compares cross-entropy and Dice loss, calibrates FCNs with model ensembling, and evaluates predictive uncertainty using pixel probabilities and average segment entropy.
Results
Across three segmentation tasks, cross-entropy improves uncertainty estimation for single models while Dice loss improves segmentation quality; ensembles with M ≥5 notably calibrate both losses, and entropy correlates with Dice quality.
Takeaways & Limitations
Model ensembling can calibrate Dice-loss FCNs, while average segment entropy can estimate segmentation performance when ground truth is unavailable and may help detect out-of-distribution samples.
Takeaways & Limitations
The segment-level uncertainty measure assumes binary classification and treats all non-foreground classes as background, omitting neighborhood relationships between classes and adjacent pixels.

Abstract

from arXiv · show

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.

I. INTRODUCTION

FCNs achieve strong medical-image segmentation performance, but their probability estimates can be overconfident and poorly calibrated. This paper studies calibration, predictive uncertainty, and out-of-distribution detection, emphasizing the risks of domain shift in medical imaging.

FCNs and U-Nets are widely used for medical image segmentation, but few studies have examined their predictive uncertainty.
Poor calibration means predicted class probabilities do not reliably correspond to the long-run probability of correct classification.
Dice loss improves robustness to class imbalance and segmentation optimization, but BN and Dice loss have been reported to harm calibration quality.
Detecting out-of-distribution inputs matters because medical-image networks are sensitive to recurring domain shifts between imaging settings.
Figure 1 contrasts uncalibrated and ensemble-calibrated prostate predictions for phased-array and endorectal-coil images, with wider probability distributions marking the out-of-distribution case.

II. RELATED WORKS

Prior work includes Bayesian uncertainty modeling and non-Bayesian calibration methods for neural networks. Ensembling and MC dropout provide alternative approaches for estimating uncertainty and improving confidence calibration.

Bayesian neural networks model predictive uncertainty by placing probability distributions over network parameters and output label probabilities.
MC dropout approximates Bayesian inference by running a network repeatedly with dropout active during inference.
Batch normalization has also been studied as an approximation to Bayesian inference for networks without dropout.
Non-Bayesian work has examined temperature scaling, network properties, and ensembling for confidence calibration and segmentation performance.

III. CONTRIBUTIONS

The paper compares cross-entropy and Dice loss, proposes ensembling for calibration, and introduces average segment entropy for predicting structure quality and detecting out-of-distribution inputs. These ideas are evaluated across brain, heart, and prostate MRI segmentation tasks.

Dice-loss FCNs achieve better segmentation than cross-entropy models, but with poorer calibration.
Model ensembles of randomly initialized and shuffled FCNs improve both segmentation quality and uncertainty estimation, and are compared with MC dropout.
Average entropy over a predicted object is proposed to predict foreground segmentation quality and detect out-of-distribution inputs.
The approach is evaluated on MRI segmentation tasks involving the brain, heart, and prostate, with statistical significance reported where appropriate.

IV. APPLICATIONS & DATA

The study uses publicly available datasets for brain tumor, heart ventricle, and prostate segmentation, with task-specific imaging modalities and data characteristics. The table organizes patient counts across training, validation, and test sets.

Table I reports the numbers of patient images assigned to training, validation, and test sets.
A separate dataset is used only for out-of-distribution detection experiments.
A. Brain Tumor Segmentation Task: Brain tumor segmentation uses four-class multiparametric MRI from the BraTS challenge, with CBICA data for training and validation and TCIA data for testing.
B. Ventricular Segmentation Task: Heart ventricle segmentation uses four-class cine MRI from the ACDC challenge and includes only end-diastole images from 100 patients.

C. Prostate Segmentation Task

The prostate task segments the prostate gland from axial T2-weighted MRI images using PROSTATEx and PROMISE12 datasets, with standardized image preprocessing.

PROSTATEx contains 40 annotated images acquired at one institution for training and testing, while PROMISE12 serves as a target dataset.
The task is binary segmentation of the prostate gland and background in axial T2-weighted MRI images.
Prostate images were resampled to 0.5×0.5 mm in-plane resolution, cropped to 224 × 224 pixels, and intensity-normalized to [0,1].

V. METHODS

The methods compare cross-entropy and Dice training, evaluate probabilistic calibration with several metrics, and use ensembles, entropy, and alternative inference procedures to quantify uncertainty.

FCNs map each input pixel to an output segmentation map, enabling end-to-end learning and direct implementation of segment-level losses such as Dice loss.
Cross-entropy optimizes average pixel-level negative log-likelihood, whereas Dice loss directly targets the Dice coefficient and is expected to improve segmentation quality.
Calibration is assessed with NLL, Brier score, reliability diagrams, and Expected Calibration Error, which summarizes accuracy–confidence gaps across bins.
Ensembles average predictions from FCNs trained with random parameter initialization and shuffled mini-batches, while MC dropout uses 50 stochastic samples at inference.
Average pixel-wise entropy over each predicted foreground segment is used as a scalar volume-level confidence metric for estimating segmentation quality and detecting difficult or out-of-distribution cases.
Experiments compare calibration and segmentation quality across three tasks using NLL, Brier score, ECE%, Dice, and 95th-percentile Hausdorff distance.

VII. RESULTS

Across prostate, heart, and brain tumor segmentation, Dice-trained ensembles improve segmentation and calibration, while segment-level entropy strongly tracks segmentation quality and identifies domain-shifted cases.

Across all tasks and segments, Dice-trained baselines outperform cross-entropy baselines in segmentation, while Dice-trained ensembles outperform all other models.
Single cross-entropy models have significantly better NLL and ECE% calibration than single Dice-trained models, while MC dropout improves calibration for cross-entropy only in prostate models.
For Dice-trained ensembles, five models reduce NLL by about 66% for prostate, 44% for heart, and 62% for brain tumor segmentation.
NLL calibration improves as ensemble size increases, and a Dice-trained ensemble with M=10 outperforms a cross-entropy baseline with M=1 for every task.
Segment-level entropy correlates strongly with logit Dice across all tasks, with correlation coefficients from 0.77 to 0.92.
In prostate testing, source-domain PROSTATEx and target-domain PROMISE12 cases cluster separately, and poorly segmented cases identified by entropy are often out-of-distribution endorectal-coil images.

VIII. DISCUSSION

Across three MRI segmentation tasks, the study finds that ensembling calibrates FCNs and that average segment entropy predicts segmentation quality and may flag out-of-distribution cases. The discussion also identifies limitations involving training cost, MRI-only evaluation, binary entropy assumptions, and the interpretation of Dice loss as a scoring rule.

Calibration and segmentation quality: Ensembling with M ≥5 notably calibrates FCNs trained with Dice loss and cross-entropy loss while also improving Dice coefficient and Hausdorff distance.The study reports reduced NLL alongside improved segmentation accuracy.
Predictive uncertainty and distribution shift: Average segment entropy predicts segment-level Dice quality at test time, and preliminary results suggest calibrated FCNs can detect out-of-distribution samples.For prostate segmentation, the ensemble correctly predicted failures associated with different imaging parameters such as imaging coils.
Loss functions and scoring rules: The authors argue that Dice loss has a fundamental mismatch with proper scoring rules because segmenters output collections of label probabilities rather than distributions over binary segmentation vectors.A binary segmentation outcome has 2^n possible vectors, whereas the network provides n marginal probabilities.
Experimental scope: The experiments covered brain, heart, and prostate MRI segmentation tasks with differences in input channels, training-set size, target-segment size, and sample homogeneity.Only publicly available datasets were used.
Loss-function comparisons: The paper compares calibration qualities of cross-entropy and Dice loss and suggests investigating combined losses and Lovász-Softmax in future studies.The discussion emphasizes that loss function choice affects both calibration quality and segmentation performance.
Metric assumptions: The proposed segment-level metric assumes binary classification and treats every non-foreground class as background when calculating foreground entropy.The authors note that multiclass entropy and neighborhood relationships could be incorporated in future work.
Limitations and future work: Calibration methods that avoid retraining ensembles remain needed because training from scratch is time-consuming, and this study evaluated uncertainty only for MR images.The authors specifically identify CT parameter changes such as slice thickness and contrast as an important future setting.

APPENDIX A CALIBRATION QUALITY (WHOLE VOLUME RESULTS)

Whole-volume calibration quality is compared across Dice-loss and cross-entropy baselines, plus ensemble and MC-dropout calibration methods.

The table reports calibration quality for each application and model configuration.
Table I compares Dice-loss baselines (LDSC) with cross-entropy baselines (LCE), ensembles with M=50, and MC dropout.
Bold entries indicate the best results for each application and model, with statistically significant differences.

APPENDIX B HAUSDORFF DISTANCE METRIC

Segmentation performance is evaluated using 95th Hausdorff distance in millimeters across brain, heart, and prostate applications.

Table II compares 95th Hausdorff distances for Dice-loss baselines, cross-entropy baselines, ensembles with M=50, and MC dropout.
Brain segments represent non-enhancing tumor, edema, and enhancing tumor, while heart segments represent the right ventricle, myocardium, and left ventricle.
The prostate application evaluates the prostate gland as its single segment.
Bold entries identify the best results for each application and model, with statistically significant differences.

APPENDIX C QUANTITATIVE RESULTS

The quantitative-results figures visualize uncertainty estimation quality for brain tumor, heart, and prostate segmentation using class probabilities and reliability diagrams.

Figure 1 compares uncertainty estimation methods for brain tumor segmentation using MRI overlays, class probabilities, and reliability diagrams.
The reliability diagrams report ECE%, NLL, and Brier score for the depicted volume.
Only reliability-diagram bins containing more than 1000 samples are shown.
Figures 2 and 3 provide corresponding example cases for heart and prostate applications.

APPENDIX D NUMBER OF MODELS IN ENSEMBLE

The paper examines ensemble size, qualitative calibration changes, segment-level uncertainty, and limited 3D-FCN generalizability across the evaluated applications.

Ensemble size: Figure 4 quantifies calibration improvement as ensemble size M increases for prostate, heart, and brain tumor segmentation, with 0.95 CI.
Qualitative ensemble effects: Figures 5 and 6 qualitatively illustrate calibration and segmentation changes for Dice-loss and cross-entropy models as M varies.
Confidence examples: Figure 7 shows predictions with different confidence levels for brain and heart applications.
Segment-level uncertainty: For ensemble-calibrated CE models with M=50, average foreground entropy shows a reverse correlation with Dice scores.
3D-FCN generalizability: Limited experiments test generalizability to 3D CNNs for prostate gland segmentation using 112 × 112 × 112-pixel inputs.
3D-FCN results: Table III reports observed average calibration quality and segmentation performance for 3D FCNs trained for prostate gland segmentation.