Unsupervised Deep-Learning based Deformable Image Registration: A Bayesian Framework https://arxiv.org/abs/2108.06771

Samah Khawaled and Moti Freiman

Abstract

Unsupervised deep-learning (DL) models were recently proposed for deformable image registration tasks. In such models, a neural-network is trained to predict the best deformation field by minimizing some dissimilarity function between the moving and the target images. After training on a dataset without reference deformation fields available, such a model can be used to rapidly predict the deformation field between newly seen moving and target images. Currently, the training process effectively provides a point-estimate of the network weights rather than characterizing their entire posterior distribution. This may result in a potential over-fitting, which may yield sub-optimal results at inference phase, especially for small-size datasets frequently present in the medical imaging domain. We introduce a fully Bayesian framework for unsupervised DL-based deformable image registration. Our method provides a principled way to characterize the true posterior distribution, thus avoiding potential over-fitting. We used stochastic gradient Langevin dynamics (SGLD) to conduct the posterior sampling, which is both theoretically well-founded and computationally efficient. We demonstrated the added-value of our Basyesian unsupervised DL-based registration framework on the MNIST and brain MRI (MGH10) datasets in comparison to the VoxelMorph unsupervised DL-based image registration framework. Our experiments show that our approach provides better estimates of the deformation field by means of improved mean-squared-error (0:0063 vs. 0:0065) and Dice coefficient (0:73 vs. 0:71) for the MNIST and the MGH10 datasets respectively. Further, our approach provides an estimate of the uncertainty in the deformation-field by characterizing the true posterior distribution.

Registration System

The following block diagram illustrates the design of our Bayesian registration framework. Our main building-block is a UNet-based CNN similar to the VoxelMorph model [1]. The operation of the system at the inference stage is as follows: it takes a pair of moving (I_M) and target (I_F) images as a 2-channel input and predicts the posterior deformation field,ϕ by computing the average of the deformation field predictions obtained by the stochastic UNets. Lastly, it maps each pixel in the moving image by applying the spatial transform function.

Experiments on 2D Examples

We implement our method for unsupervised 2D deformable image registration on both the MNIST [2] and the MGH10 brain MRI [3] datasets.

Samah presenting her research for the ISMRM 2021

References

[1] Guha Balakrishnan, Amy Zhao, Mert R Sabuncu, John Guttag, and Adrian V Dalca. An unsupervised learning model for deformable medical image registration. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 9252–9260, 2018.
[2] Yann LeCun. The mnist database of handwritten digits. yann. lecun. com/exdb/mnist/, 1998.
[3] Arno Klein, Jesper Andersson, Babak A Ardekani, John Ashburner, Brian Avants, Ming-Chang Chiang, Gary E Christensen, D Louis Collins, James Gee, Pierre Hellier, et al. Evaluation of 14 nonlinear deformation algorithms applied to human brain mri registration. Neuroimage, 46(3):786–802, 2009.