Uncertainty quantification: estimating aleatoric and epistemic uncertainty in medical image segmentation

Abstract

Medical imaging is a cornerstone for medical diagnosis, treatment planning, and clinical studies. In order to delineate anatomical structures and other regions of interest in such images, deep neural networks can be employed, performing image segmentation for the medical expert. Because of the high-risk setting, these models need to be not only exact and robust, but also indicate error likelihood via reliable and meaningful uncertainty estimates. This predictive uncertainty can be subdivided into aleatoric and epistemic uncertainty, and captured using deep ensembles, Bayesian neural networks, and additional loss attenuating output neurons. The main contribution of this work is a comprehensive comparison between the direct decomposition of uncertainty in Bayesian neural networks via the mutual information metric and the explicit modelling of epistemic and aleatoric uncertainty in Bayesian neural networks with an additional heteroscedastic loss-attenuating neuron. This comparison is performed in the context of medical image segmentation of the liver from CT scans, employing a 3D au-net as base architecture. The quality of the uncertainty decomposition in the resulting uncertainty maps is qualitatively evaluated and quantitative behaviour of aleatoric and epistemic uncertainty is systematically compared for different experiment settings with varying training set sizes, label noise, and distribution shifts. The results show the mutual information decomposition to robustly yield meaningful aleatoric and epistemic uncertainty estimates, with both largely conforming to their definitions and consistent with other works. Noisiness in the activation of the loss-attenuating neuron leads to the conclusion that the mutual information decomposition remains significantly more suited for uncertainty decomposition even for Bayesian neural networks combined with loss-attenuating neurons. This work further found that the addition of a heteroscedastic neuron does not improve the quality of the uncertainty estimates when decomposed via the mutual information metric. An ancillary contribution is the demonstration of a strong influence of the choice of loss function on the quality of uncertainty decomposition, with soft Dice loss heavily deteriorating the quality of the decomposed uncertainties.