Beyond scalar losses: calibrating segmentation models via gradient vector field surgery

Laurin Lux, Alexander H Berger, Moritz Knolle, Daniel Rueckert, Johannes C. Paetzold

Keywords: calibration, image segmentation, gradient surgery, tumor segmentation

TL;DR: Combating miscalibration in segmentation networks with a "gradient surgery" approach that modifies the gradient dynamics of region-based loss functions to scale linearly with probabilistic error.

Abstract: Region-based loss functions, such as the Dice loss, have established themselves as the de facto standard for highly class- and region-imbalanced segmentation tasks. However, models trained using region-based loss functions are notoriously miscalibrated and typically yield over-confident predictions. In medical imaging applications, such as defining tumor resection margins, this miscalibration is hindering clinical adoption. In this work, we outline a novel gradient perspective on this overconfidence and show how it affects region-based loss functions. We propose a "surgery" on the gradient vector field as a simple, yet effective intervention to mitigate calibration issues. This surgery adds a factor to the loss's partial derivative, scaling the gradient's magnitude linearly with the prediction error. In empirical evaluations across 2D and 3D medical segmentation tasks, we demonstrate how this intervention improves model calibration while maintaining high prediction accuracy when used in conjunction with any region-based loss function.

Primary Subject Area: Segmentation

Secondary Subject Area: Safe and Trustworthy Learning-assisted Solutions for Medical Imaging

Registration Requirement: Yes

Visa & Travel: No

Read CFP & Author Instructions: Yes

Originality Policy: Yes

Single-blind & Not Under Review Elsewhere: Yes

LLM Policy: Yes

Submission Number: 317