Efficient Bayesian Uncertainty Estimation for nnU-Net

  • 2024-05-01 07:49:03
  • Yidong Zhao, Changchun Yang, Artur Schweidtmann, Qian Tao
  • 0

Abstract

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

 

Quick Read (beta)

loading the full paper ...