Learning a Probabilistic Model for Diffeomorphic Registration

Abstract

We propose to learn a low-dimensional probabilistic deformation model from data which can be used for registration and the analysis of deformations. The latent variable model maps similar deformations close to each other in an encoding space. It enables to compare deformations, generate normal or pathological deformations for any new image or to transport deformations from one image pair to any other image.Our unsupervised method is based on variational inference. In particular, we use a conditional variational autoencoder (CVAE) network and constrain transformations to be symmetric and diffeomorphic by applying a differentiable exponentiation layer with a symmetric loss function. We also present a formulation that includes spatial regularization such as diffusion-based filters. Additionally, our framework provides multi-scale velocity field estimations. We evaluated our method on 3-D intra-subject registration using 334 cardiac cine-MRIs. On this dataset, our method showed state-of-the-art performance with a mean DICE score of 81.2% and a mean Hausdorff distance of 7.3mm using 32 latent dimensions compared to three state-of-the-art methods while also demonstrating more regular deformation fields. The average time per registration was 0.32s. Besides, we visualized the learned latent space and show that the encoded deformations can be used to transport deformations and to cluster diseases with a classification accuracy of 83% after applying a linear projection.

Probabilistic Motion Modeling from Medical Image Sequences: Application to Cardiac Cine-MRI

“…The motion observed in an image sequence with T + 1 frames is typically described by deformation fields φ t between a moving image I 0 and the fixed images I t with t ∈ [1, T ]. Inspired by the probabilistic deformation model of [10] based on conditional variational autoencoder (CVAE) [9], we define a motion model for temporal sequences. The model is conditioned on the moving image and parameterizes the set of diffeomorphisms φ t in a low-dimensional probabilistic space, the motion matrix z ∈ R d×T , where d is the size of the deformation encoding per image pair.…”

“…We model p θ as a symmetric local cross-correlation Boltzmann distribution with the weighting factor λ. Encoder and decoder weights are independent of the time t. Their network architecture consists of convolutional and deconvolutional layers with fully-connected layers for mean and variance predictions in the encoder part [9]. We use an exponentiation layer for the stationary velocity field parameterization of diffeomorphisms [10], a linear warping layer and diffusion-like regularization with smoothing parameters σ G in spatial and σ T in temporal dimension.…”

“…LV volumes in ml are shown for two test sequences (ground truth ED/ES volumes marked with points). The proposed algorithms (Our and Our w/o TDS) show slightly higher registration accuracy and temporally smoother deformations than the state-of-the-art algorithms: SyN[2], LPR[10] and 4D-Elastix[12].…”

Unsupervised Deep Learning for Bayesian Brain MRI Segmentation

“…From this point, an additional convolutional layer is used to output v, a dense 3D velocity field defined over Ω; and an additional pair of convolutional layers followed by a global max pooling operation to output the Gaussian parameters µ, σ 2 . We compute φ = exp(v) using a network integration layer that implements scaling and squaring [2,8,21], enabling the computation of the loss regularization term. We warp the probabilistic atlas A with a spatial transform layer.…”

“…The estimation of the atlas warp has traditionally relied on classic deformable registration algorithms [31], which are based on iterative, numerical optimization, and are therefore computationally expensive. Instead, we leverage recent advances in learning-based registration [4,8,21,33] to efficiently estimate the warp jointly with the intensity parameters. We use a novel loss function, which is derived from the probabilistic model with Bayesian inference, and is thus principled and interpretable.…”