Non-parametric Uncertainty Assessment in Deep-Learning Based Affine Image Registration

Samah Khawaled, Ameer Abu Hanna, Hussain Atieh, Moti Freiman

Keywords: Affine Registration, Brain MRI, Uncertainty Assessment

TL;DR: Uncertainty Assessment in Deep-Learning Based Affine Image Registration for Brain MRI Images

Abstract:

Affine image registration plays a key role in diagnosis, surgical planning and in data-processing pipelines for research as both an essential initialization for subsequent non-rigid registration or as an independent step. Uncertainty quantification in deep learning (DL) based image registration models is critical for determining confidence intervals required for surgical guidance, and for reliable assessment of differences between the registered images.
We introduce AIR-SGLD - a non-parametric fully Bayesian framework for Affine image registration. We use Stochastic gradient Langevin dynamics (SGLD) during the training phase to characterize the posterior distribution of the network weights.
We demonstrated the added-value of AIR-SGLD on the brain MRI (MGH10) dataset in comparison to the baseline AIR DL-based Affine image registration framework using 300 pairs of images generated from the MGH10 dataset. Our experiments show that AIR-SGLD outperforms AIR by means of cross-correlation between the images ($0.91$ vs. $0.87$, $p<0.01$). Further, AIR-SGLD provides an estimate of the registration uncertainty that correlates with both registration error (Pearson correlation coefficient of $R=0.769$) and the presence of out-of-distribution data ($R=0.796$).
AIR-SGLD has the potential to provide reliable and more accurate registration for clinical diagnosis, surgical planning, and automatic data processing pipelines.