Assessing Uncertainty Estimation Methods for 3D Image Segmentation under Distribution Shifts

TL;DR

The paper addresses the critical issue of reliable uncertainty estimation for 3D medical image segmentation under distribution shifts. It compares multimodal-posterior capable methods (cSGHMC) with non-Bayesian and single-mode approaches (MCD, DE) across covariate, modality, and corruption shifts using three 3D datasets, evaluating calibration and predictive uncertainty via $NLL$, $BS$, and $ECE$ as well as entropy-based measures. Results indicate that methods capturing multiple posterior modes yield more trustworthy uncertainty estimates, with cSGHMC often providing superior calibration and higher, more informative uncertainty under shift, while Deep Ensemble shows limited diversity and inconsistent reliability. The study offers practical guidance for deploying medical AI with uncertainty thresholds to avoid making dangerous decisions when faced with distribution shifts, underscoring the importance of posterior diversity in robust uncertainty quantification.

Abstract

In recent years, machine learning has witnessed extensive adoption across various sectors, yet its application in medical image-based disease detection and diagnosis remains challenging due to distribution shifts in real-world data. In practical settings, deployed models encounter samples that differ significantly from the training dataset, especially in the health domain, leading to potential performance issues. This limitation hinders the expressiveness and reliability of deep learning models in health applications. Thus, it becomes crucial to identify methods capable of producing reliable uncertainty estimation in the context of distribution shifts in the health sector. In this paper, we explore the feasibility of using cutting-edge Bayesian and non-Bayesian methods to detect distributionally shifted samples, aiming to achieve reliable and trustworthy diagnostic predictions in segmentation task. Specifically, we compare three distinct uncertainty estimation methods, each designed to capture either unimodal or multimodal aspects in the posterior distribution. Our findings demonstrate that methods capable of addressing multimodal characteristics in the posterior distribution, offer more dependable uncertainty estimates. This research contributes to enhancing the utility of deep learning in healthcare, making diagnostic predictions more robust and trustworthy.

Figures (13)

• Figure 1: Performance and Calibration Metrics Comparison on Blurred Images. The first row indicates performance and various calibration metrics. The second row presents the histogram of predictive uncertainty across different models for shift equal to 6px, with cSGHMC and MCD exhibiting the highest uncertainty.
• Figure 2: Performance and Calibration Metrics Comparison on Rotated Images. The first row presents a comparison of performance and various calibration metrics. The second row illustrates uncertainty estimation across different models for shift equal to $30^\circ$, with cSGHMC exhibiting the highest uncertainty
• Figure 3: Performance and Calibration Metrics Comparison on AMOS dataset. First row includes, reliability diagram, ECE and performance metrics. The second row depicts the histograms of predictive uncertainty across various methods, with DE and cSGHMC exhibiting slightly higher uncertainty.
• Figure 4: Performance and Calibration Metrics Comparison on KITS dataset. First row encompasses, reliability diagram, ECE and performance metrics. The second row depicts the histograms of predictive uncertainty across different methods, highlighting cSGHMC as the most calibrated with the highest uncertainty.
• Figure 5: Pairwise correlation of softmax outputs between any two model samples in the posterior distribution for cSGHMC and in the predictive space for MCD and DE. For cSGHMC, samples from the last 6 cycles were chosen. For MCD and DE, 6 models were randomly selected.