Uncertainty Estimation and Calibration in nnU-Net-Based Ischemic Stroke Lesion Segmentation

Ewout Heylen, Jelle Demeestere, Anke Wouters, Pierre Seners, Soren Christensen, Nicole Yuen, Michael Mlynash, Stephanie Kemp, Jeremy J. Heit, Gregory W. Albers, Seena Dehkharghani, Robin Lemmens, Maarten Lansberg, Frederik Maes

Keywords: nnU-Net, Model Calibration, Uncertainty Estimation, Ischemic Stroke

TL;DR: Bayesian posterior sampling for uncertainty estimation and calibration in nnU-Net.

Abstract: nnU-Net has become widely recognized as a state-of-the-art semantic segmentation framework. However, deep learning models are often poorly calibrated, resulting in unreliable probability estimates. Additionally, they lack meaningful uncertainty quantification. We trained an nnU-Net model to segment ischemic stroke lesions on acute-phase Diffusion-Weighted Imaging (DWI) MRI and applied a Bayesian posterior sampling approach to estimate uncertainty and improve model calibration. Our findings show that the Bayesian posterior sampling approach yields better calibration compared to a conventional nnU-Net, while providing uncertainty estimates and maintaining comparable segmentation performance.

Submission Number: 110