Bayesian Uncertainty Estimation by Hamiltonian Monte Carlo: Applications to Cardiac MRI Segmentation

Yidong Zhao10000-0003-3953-6921, João Tourais10000-0002-1388-4023, Iain Pierce2,3, Christian Nitsche2,3, Thomas A. Treibel2,3, Sebastian Weingärtner10000-0002-0739-6306, Artur M. Schweidtmann40000-0001-8885-6847, Qian Tao10000-0001-7480-0703
1: Department of Imaging Physics, Delft University of Technology, Delft, The Netherlands, 2: Barts Heart Centre, Barts Health NHS Trust, London, United Kingdom, 3: Institute of Cardiovascular Science, University College, London, United Kingdom, 4: Department of Chemical Engineering, Delft University of Technology, Delft, The Netherlands
Publication date: 2024/06/23
https://doi.org/10.59275/j.melba.2024-88fa
PDF · Code · arXiv

Abstract

Deep learning (DL)-based methods have achieved state-of-the-art performance for many medical image segmentation tasks. Nevertheless, recent studies show that deep neural net- works (DNNs) can be miscalibrated and overconfident, leading to ”silent failures” that are risky for clinical applications. Bayesian DL provides an intuitive approach to DL failure de- tection, based on posterior probability estimation. However, the posterior is intractable for large medical image segmentation DNNs. To tackle this challenge, we propose a Bayesian learning framework using Hamiltonian Monte Carlo (HMC), tempered by cold posterior (CP) to accommodate medical data augmentation, named HMC-CP. For HMC compu- tation, we further propose a cyclical annealing strategy, capturing both local and global geometries of the posterior distribution, enabling highly efficient Bayesian DNN training with the same computational budget as training a single DNN. The resulting Bayesian DNN outputs an ensemble segmentation along with the segmentation uncertainty. We evaluate the proposed HMC-CP extensively on cardiac magnetic resonance image (MRI) segmentation, using in-domain steady-state free precession (SSFP) cine images as well as out-of-domain datasets of quantitative T1 and T2 mapping. Our results show that the proposed method improves both segmentation accuracy and uncertainty estimation for in- and out-of-domain data, compared with well-established baseline methods such as Monte Carlo Dropout and Deep Ensembles. Additionally, we establish a conceptual link between HMC and the commonly known stochastic gradient descent (SGD) and provide general insight into the uncertainty of DL. This uncertainty is implicitly encoded in the training dynamics but often overlooked. With reliable uncertainty estimation, our method provides a promising direction toward trustworthy DL in clinical applications. We release our code in https://gitlab.tudelft.nl/yidongzhao/hmc_uncertainty

Keywords

Uncertainty estimation · Bayesian deep learning · Hamiltonian Monte Carlo · segmentation · cardiac MRI

Bibtex @article{melba:2024:011:zhao, title = "Bayesian Uncertainty Estimation by Hamiltonian Monte Carlo: Applications to Cardiac MRI Segmentation", author = "Zhao, Yidong and Tourais, João and Pierce, Iain and Nitsche, Christian and Treibel, Thomas A. and Weingärtner, Sebastian and Schweidtmann, Artur M. and Tao, Qian", journal = "Machine Learning for Biomedical Imaging", volume = "2", issue = "June 2024 issue", year = "2024", pages = "856--887", issn = "2766-905X", doi = "https://doi.org/10.59275/j.melba.2024-88fa", url = "https://melba-journal.org/2024:011" }
RISTY - JOUR AU - Zhao, Yidong AU - Tourais, João AU - Pierce, Iain AU - Nitsche, Christian AU - Treibel, Thomas A. AU - Weingärtner, Sebastian AU - Schweidtmann, Artur M. AU - Tao, Qian PY - 2024 TI - Bayesian Uncertainty Estimation by Hamiltonian Monte Carlo: Applications to Cardiac MRI Segmentation T2 - Machine Learning for Biomedical Imaging VL - 2 IS - June 2024 issue SP - 856 EP - 887 SN - 2766-905X DO - https://doi.org/10.59275/j.melba.2024-88fa UR - https://melba-journal.org/2024:011 ER -

2024:011 cover