Efficient Bayesian Uncertainty Estimation for nnU-Net

TL;DR

This paper introduces a novel Bayesian uncertainty estimation method for nnU-Net that preserves its architecture and improves both segmentation accuracy and uncertainty quantification in medical imaging.

Contribution

The authors propose a new posterior sampling scheme for Bayesian uncertainty estimation that does not alter nnU-Net's architecture, enhancing its performance and reliability.

Findings

Improved uncertainty estimation over baseline methods.

Enhanced segmentation accuracy and quality control.

Validated on cardiac MRI datasets with superior results.

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.