Efficient Bayesian Uncertainty Estimation for nnU-Net

TL;DR

This paper addresses the lack of uncertainty estimation in the self-configuring nnU-Net for medical image segmentation by introducing a trajectory-based Bayesian approach that preserves the original architecture. It leverages SGD weight-space sampling to approximate the posterior $p(\mathbf{w}|\mathcal{D})$ and uses both single-modal and multi-modal posterior sampling, including a cyclical learning-rate strategy and multi-cycle checkpoint ensembles, to improve predictive uncertainty and calibration. Experiments on cardiac MRI datasets (ACDC for ID and M&M for OOD) show that the proposed method improves segmentation performance and calibration compared to MC-Dropout and Deep Ensemble, with multi-modal ensembles particularly enhancing OOD robustness. The approach yields uncertainty maps that correlate with hard regions and potential failures, supporting improved quality control in large-scale deployments without modifying nnU-Net’s architecture.

Abstract

The self-configuring nnU-Net has achieved leading performance in a large range of medical image segmentation challenges. It is widely considered as the model of choice and a strong baseline for medical image segmentation. However, despite its extraordinary performance, nnU-Net does not supply a measure of uncertainty to indicate its possible failure. This can be problematic for large-scale image segmentation applications, where data are heterogeneous and nnU-Net may fail without notice. In this work, we introduce a novel method to estimate nnU-Net uncertainty for medical image segmentation. We propose a highly effective scheme for posterior sampling of weight space for Bayesian uncertainty estimation. Different from previous baseline methods such as Monte Carlo Dropout and mean-field Bayesian Neural Networks, our proposed method does not require a variational architecture and keeps the original nnU-Net architecture intact, thereby preserving its excellent performance and ease of use. Additionally, we boost the segmentation performance over the original nnU-Net via marginalizing multi-modal posterior models. We applied our method on the public ACDC and M&M datasets of cardiac MRI and demonstrated improved uncertainty estimation over a range of baseline methods. The proposed method further strengthens nnU-Net for medical image segmentation in terms of both segmentation accuracy and quality control.

Figures (3)

• Figure 1: We observe that the network checkpoints at various training epochs make diverse predictions when the network is uncertain on a hard input $\mathbf{x}^h$ (a). On the contrary, predictions of good quality on an easy input $\mathbf{x}^e$ enjoys consistency across checkpoints (b). We leverage this phenomenon to perform Bayesian inference and quantify uncertainty of network predictions.
• Figure 2: (a) t-SNE plot of the weight space during SGD training. Dotted lines illustrate the transition between weight modes. (b) t-SNE plot of the posterior weights, which bounce in the neighborhood of different modes.
• Figure 3: Predictions (Pred.) and estimated uncertainty maps (Uncert.) on a successful case (a) and three partially failed cases (b-d). In case (a) all the methods highlighted the border as uncertain. MC-Dropout failed in all the three cases (b-d) to delineate RV, while reporting low uncertainty in the corresponding area. Deep Ensemble is robust but missed part of the uncertain areas in case (c). Multi-modal weight sampling detected the failed RV area more robustly than the single-modal version in case (d).