Average Calibration Losses for Reliable Uncertainty in Medical Image Segmentation

1. Problem Statement

Deep Neural Networks (DNNs) used for medical image segmentation are frequently overconfident, meaning their predicted probabilities (confidence) do not align with their actual accuracy. In safety-critical clinical domains, this miscalibration poses significant risks, such as misleading clinicians with high-confidence incorrect predictions.

Existing solutions face limitations:

• Post-hoc methods (e.g., Temperature Scaling, Platt Scaling) apply corrections after training. They often fail to leverage the model's full capacity to learn calibration and rely on global adjustments that cannot correct instance- or class-specific errors.
• Standard training losses (e.g., Cross-Entropy, Dice Loss) do not explicitly optimize for calibration. While Cross-Entropy is a proper scoring rule, it often yields overconfident predictions in segmentation tasks.
• Calibration Metrics: The widely used Expected Calibration Error (ECE) weights errors by bin population, making it dominated by highly confident predictions (often the background class) and insensitive to errors in critical, uncertain regions (e.g., tissue boundaries).

2. Methodology

The authors propose integrating calibration directly into the training process using a differentiable auxiliary loss function based on the Marginal L1 Average Calibration Error (mL1-ACE).

Core Concepts

• Per-Image Calibration: Unlike classification tasks where calibration is computed over a dataset, segmentation outputs dense voxel-level predictions. The authors exploit this to compute reliability diagrams and calibration metrics on a per-image basis, providing stable estimates even from single images.
• Metric Choice (ACE vs. ECE): The method uses Average Calibration Error (ACE) instead of ECE. ACE treats all confidence bins equally, whereas ECE weights bins by sample count. This ensures that errors in low-confidence or boundary regions (clinically critical) are not drowned out by the high volume of confident background predictions.
• Differentiable Formulation: To make ACE trainable, the authors define two binning strategies to compute the loss:
  1. Hard-Binned mL1-ACE (hL1-ACE): Uses a square kernel (step function) to assign voxels to discrete bins. This preserves sharp decision boundaries.
  2. Soft-Binned mL1-ACE (sL1-ACE): Uses a triangular kernel (continuous overlap) where voxels contribute to multiple adjacent bins. This creates smoother gradients and stronger regularization of probabilities.

Loss Function

The total loss function combines the standard segmentation loss with the calibration loss:
$$\mathcal{L}_{total} = \mathcal{L}_{DSC} + \mathcal{L}_{CE} + \lambda \cdot \mathcal{L}_{mL1-ACE}$$
Where $\mathcal{L}_{DSC}$ is the Dice Similarity Coefficient loss, $\mathcal{L}_{CE}$ is Cross-Entropy, and $\lambda$ is a weighting hyperparameter.

New Visualization Tool

The paper introduces Dataset Reliability Histograms. Instead of a single reliability diagram for the whole dataset, this tool aggregates per-image reliability diagrams. It visualizes the distribution of observed accuracies ($o_m$) for each predicted confidence bin ($e_m$) across the entire dataset, offering a granular view of calibration variability.

3. Key Contributions

1. Differentiable mL1-ACE Loss: The first formulation of marginal L1 Average Calibration Error as a differentiable auxiliary loss specifically for semantic segmentation, computable on a per-image basis.
2. Binning Strategy Comparison: A systematic comparison of hard-binning (discrete) vs. soft-binning (continuous) approaches, revealing a distinct trade-off between calibration improvement and segmentation accuracy.
3. Dataset Reliability Histograms: A novel visualization technique to analyze calibration performance and variability across large medical imaging datasets.
4. Comprehensive Evaluation: Extensive experiments on four diverse public datasets (ACDC, AMOS, BraTS, KiTS) covering cardiac MRI, abdominal CT, brain tumor MRI, and kidney CT.

4. Experimental Results

The models were evaluated against a DSC + CE baseline, a CE-only baseline, and Post-hoc Temperature Scaling.

• Calibration Performance:
  • Both hL1-ACE and sL1-ACE significantly reduced calibration errors (ACE and MCE) compared to the baseline.
  • Soft-binned (sL1-ACE) achieved the greatest calibration improvements, reducing macro-averaged ACE by 33% and micro-averaged ACE by 51% on average across datasets.
  • Hard-binned (hL1-ACE) provided moderate calibration gains (16% macro-ACE reduction) but maintained segmentation performance almost perfectly.
• Segmentation Performance (DSC):
  • hL1-ACE: Maintained DSC scores nearly identical to the baseline (average change: +0.03%).
  • sL1-ACE: Achieved superior calibration but at the cost of a slight reduction in segmentation accuracy (average DSC decrease: -0.7%).
• Comparison to Baselines:
  • Both variants outperformed the CE-only baseline in calibration while preserving the high DSC of the DSC+CE baseline.
  • hL1-ACE achieved calibration comparable to Post-hoc Temperature Scaling but with the advantage of being learned during training (no separate calibration set needed).
• Sensitivity Analysis:
  • Hard-binning was robust to changes in bin count and loss weighting ($\lambda$).
  • Soft-binning was more sensitive; higher $\lambda$ and more bins led to sharper drops in DSC, indicating a stronger accuracy-calibration trade-off.

5. Significance and Implications

• Explicit Control over Trade-offs: The method provides practitioners with a tunable knob ($\lambda$ and binning type) to explicitly manage the calibration-accuracy trade-off.
  • Use Hard-binning when segmentation accuracy is paramount (e.g., surgical planning) but some calibration improvement is desired.
  • Use Soft-binning when reliability and uncertainty quantification are critical (e.g., active learning, quality control, safety-critical triage), accepting a minor drop in DSC.
• Clinical Utility: By reducing overconfidence, these models can better flag uncertain regions (e.g., tumor boundaries) for human review, potentially improving clinical decision-making and patient outcomes.
• Metric Advancement: The paper argues for the adoption of ACE over ECE in medical segmentation, as ACE better captures errors in critical, low-frequency regions that ECE often ignores.
• Limitations: The authors note that images with missing ground-truth classes (common in multi-organ datasets) can skew calibration metrics. Additionally, binning strips spatial information, which limits the modeling of spatial uncertainty correlations.

In conclusion, this work demonstrates that train-time calibration via differentiable auxiliary losses is a viable and effective strategy for producing reliable, trustworthy AI models in medical imaging, offering a flexible alternative to post-hoc correction methods.