Boundary-weighted logit consistency improves calibration of segmentation networks

TL;DR

This paper introduces a boundary-weighted logit consistency regularizer that improves the calibration of segmentation networks, especially in regions with ambiguous labels, leading to more reliable probability estimates.

Contribution

The authors propose a novel boundary-weighted regularizer based on logit consistency across stochastic transformations for better calibration in segmentation tasks.

Findings

Achieves state-of-the-art calibration on prostate MRI segmentation

Effectively reduces overconfidence at ambiguous boundaries

Improves reliability of segmentation probability estimates

Abstract

Neural network prediction probabilities and accuracy are often only weakly-correlated. Inherent label ambiguity in training data for image segmentation aggravates such miscalibration. We show that logit consistency across stochastic transformations acts as a spatially varying regularizer that prevents overconfident predictions at pixels with ambiguous labels. Our boundary-weighted extension of this regularizer provides state-of-the-art calibration for prostate and heart MRI segmentation.