Uncertainty quantification in non-rigid image registration via stochastic gradient Markov chain Monte Carlo

TL;DR

This paper introduces a Bayesian approach for non-rigid 3D medical image registration that efficiently quantifies uncertainty, overcoming computational and modeling challenges, and compares favorably to existing deep learning methods.

Contribution

It presents a novel Bayesian model that enables calibrated uncertainty quantification in high-dimensional, diffeomorphic image registration tasks.

Findings

Improved uncertainty calibration over existing models

Efficient sampling via connections between MCMC and variational inference

Competitive registration accuracy compared to VoxelMorph

Abstract

We develop a new Bayesian model for non-rigid registration of three-dimensional medical images, with a focus on uncertainty quantification. Probabilistic registration of large images with calibrated uncertainty estimates is difficult for both computational and modelling reasons. To address the computational issues, we explore connections between the Markov chain Monte Carlo by backpropagation and the variational inference by backpropagation frameworks, in order to efficiently draw samples from the posterior distribution of transformation parameters. To address the modelling issues, we formulate a Bayesian model for image registration that overcomes the existing barriers when using a dense, high-dimensional, and diffeomorphic transformation parametrisation. This results in improved calibration of uncertainty estimates. We compare the model in terms of both image registration accuracy and uncertainty quantification to VoxelMorph, a state-of-the-art image registration model based on deep learning.