Calibrating the Dice loss to handle neural network overconfidence for biomedical image segmentation

TL;DR

This paper introduces DSC++, an improved loss function for biomedical image segmentation that enhances calibration of neural network predictions, making them more interpretable and useful in clinical settings.

Contribution

The study proposes DSC++, a novel extension of the Dice loss that improves calibration of segmentation models, addressing a key limitation of traditional DSC loss.

Findings

DSC++ significantly improves calibration across multiple datasets.

Integration of DSC++ enhances the calibration of existing DSC-based loss functions.

Well-calibrated outputs enable better post-processing adjustments like softmax thresholding.

Abstract

The Dice similarity coefficient (DSC) is both a widely used metric and loss function for biomedical image segmentation due to its robustness to class imbalance. However, it is well known that the DSC loss is poorly calibrated, resulting in overconfident predictions that cannot be usefully interpreted in biomedical and clinical practice. Performance is often the only metric used to evaluate segmentations produced by deep neural networks, and calibration is often neglected. However, calibration is important for translation into biomedical and clinical practice, providing crucial contextual information to model predictions for interpretation by scientists and clinicians. In this study, we provide a simple yet effective extension of the DSC loss, named the DSC++ loss, that selectively modulates the penalty associated with overconfident, incorrect predictions. As a standalone loss function, the DSC++ loss achieves significantly improved calibration over the conventional DSC loss across six well-validated open-source biomedical imaging datasets, including both 2D binary and 3D multi-class segmentation tasks. Similarly, we observe significantly improved calibration when integrating the DSC++ loss into four DSC-based loss functions. Finally, we use softmax thresholding to illustrate that well calibrated outputs enable tailoring of recall-precision bias, which is an important post-processing technique to adapt the model predictions to suit the biomedical or clinical task. The DSC++ loss overcomes the major limitation of the DSC loss, providing a suitable loss function for training deep learning segmentation models for use in biomedical and clinical practice. Source code is available at: https://github.com/mlyg/DicePlusPlus.