Average Calibration Losses for Reliable Uncertainty in Medical Image Segmentation

TL;DR

The paper tackles miscalibration in medical image segmentation by introducing differentiable per-image marginal L1 calibration loss (mL1-ACE) with hard and soft binning. By integrating mL1-ACE into training, it achieves significant reductions in ACE and MCE across four datasets while largely preserving Dice scores, with soft binning delivering the strongest calibration gains at some segmentation cost. It also introduces dataset reliability histograms to visualize calibration variability and explores both micro- and macro-averaging to comprehensively evaluate reliability. Overall, the method offers a practical route to more trustworthy segmentation outputs and informs safer clinical deployment with improved confidence calibration.

Abstract

Deep neural networks for medical image segmentation are often overconfident, compromising both reliability and clinical utility. In this work, we propose differentiable formulations of marginal L1 Average Calibration Error (mL1-ACE) as an auxiliary loss that can be computed on a per-image basis. We compare both hard- and soft-binning approaches to directly improve pixel-wise calibration. Our experiments on four datasets (ACDC, AMOS, KiTS, BraTS) demonstrate that incorporating mL1-ACE significantly reduces calibration errors, particularly Average Calibration Error (ACE) and Maximum Calibration Error (MCE), while largely maintaining high Dice Similarity Coefficients (DSCs). We find that the soft-binned variant yields the greatest improvements in calibration, over the Dice plus cross-entropy loss baseline, but often compromises segmentation performance, with hard-binned mL1-ACE maintaining segmentation performance, albeit with weaker calibration improvement. To gain further insight into calibration performance and its variability across an imaging dataset, we introduce dataset reliability histograms, an aggregation of per-image reliability diagrams. The resulting analysis highlights improved alignment between predicted confidences and true accuracies. Overall, our approach not only enhances the trustworthiness of segmentation predictions but also shows potential for safer integration of deep learning methods into clinical workflows. We share our code here: https://github.com/cai4cai/Average-Calibration-Losses

Figures (6)

• Figure 1: Reliability diagrams for BraTS 2021 (case 00095, whole tumour class), comparing baseline DSC plus CE loss (left) and baseline plus our auxiliary soft-binned mL1-ACE (sL1-ACE) loss (right). Each diagram shows empirical foreground frequency (accuracy) versus predicted foreground probability (confidence), voxel counts per bin, and associated calibration errors (ACE, ECE, MCE). The additional sL1-ACE loss substantially reduces the baseline model's overconfidence.
• Figure 2: Visualisation of hard- (solid line) and soft- (dashed line) foreground probability binning assignment functions for 5 bins.
• Figure 3: Normalised radar plots showing DSC, ACE, ECE and MCE across the four testing datasets and, in the rightmost column, for the KiTS testing dataset excluding cases with missing classes, keeping only complete cases. Models were trained with three different losses. Blue: baseline (DSC + CE), Orange: baseline plus hard-binned mL1-ACE (hL1-ACE), Green: baseline plus soft-binned L1-ACE (sL1-ACE).
• Figure 4: Dataset reliability histograms averaged over classes from each dataset, for baseline loss (DSC + CE) and baseline plus mL1-ACE with hard (hL1-ACE) and soft (sL1-ACE) binning. Gamma correction is applied to the histograms for better visualisation of lower frequencies.
• Figure 5: Comparison of predicted segmentation probability maps between baseline and models trained with our auxiliary calibration losses. The leftmost column shows input images with ground truth segmentations. Rows represent different anatomical classes from selected dataset cases. Overlaid contours highlight prediction alignment with ground truth: true positives in blue, false negatives in red, and false positives in yellow.