Learning a Probabilistic Model for Diffeomorphic Registration

TL;DR

This paper introduces a probabilistic, diffeomorphic registration model using a variational autoencoder that learns deformation representations, enabling efficient registration, deformation analysis, and disease clustering with state-of-the-art accuracy.

Contribution

It presents a novel unsupervised CVAE-based framework for diffeomorphic registration that incorporates spatial regularization and multi-scale velocity estimation, advancing deformation modeling.

Findings

Achieved 81.2% DICE score on cardiac MRI registration

Demonstrated state-of-the-art performance with 0.32s registration time

Enabled deformation transport and disease clustering with 83% accuracy.

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.

Figures

Tables

Table 1. TABLE I: Registration performance with mean and standard deviation scores (in brackets) of RMSE, DICE, Hausdorff Distance (HD in mm ) and the mean gradient of the determinant of Jacobians (Grad Det-Jac, × 10 − 2 absent superscript 10 2 \times 10^{-2} ) comparing the undeformed case (Undef), LCC-demons (Dem), SyN , voxelmorph (VM) and our method.

Method	RMSE	DICE	HD	Grad Det-Jac
Undef	0.37 (0.17)	0.707 (0.145)	10.1 (2.2)	–
Dem	0.29 (0.16)	0.799 (0.096)	8.3 (2.7)	2.9 (1.0)
SyN	0.32 (0.16)	0.801 (0.091)	8.1 (3.6)	3.4 (0.5)
VM	0.24 (0.08)	0.790 (0.096)	8.4 (2.6)	9.2 (0.5)
Our S1	0.31 (0.15)	0.797 (0.093)	7.9 (2.6)	1.2 (0.3)
Our S3	0.30 (0.14)	0.812 (0.085)	7.3 (2.7)	1.4 (0.3)

Table 2. TABLE II: Mean Ejection fraction (EF in % with standard deviation in parentheses) of pathological deformation predictions (Step 1) should stay similar to the mean EF after the transport to healthy/normal subjects (Step 2). Our algorithm shows smaller absolute differences compared to the pole ladder (PL).

	Step 1: Prediction		Step 2: Transport		Difference
Path.	PL (Dem)	Our	PL (Dem)	Our	PL	Our
HCM	29.4 (6)	44.1 (7)	35.5 (8)	38.6 (13)	6.1	5.5
DCM	10.8 (2)	12.7 (7)	16.7 (4)	13.9 (7)	5.9	1.2

Equations

$$p(z)\sim \mathcal{N}(0,I).$$

$$v^s=f_v^s(z,M^s;{\theta }^s).$$

$$M^{\ast s}:=f^s(z,M^s;{\theta }^s)=M^s\circ \text{exp}(f_v^s(z,M^s;{\theta }^s)).$$

$$p_{\theta }(F|z,M)=\prod _{s\in \mathcal{S}}p(F^s|z,M^s;{\theta }^s).$$

$$p_{\theta }(F|M)={\int }_z{p}_{\theta }(F|z,M)p(z)dz.$$

$$p_{\theta }^s(F^s|z,M^s)\sim \text{exp}(-\lambda {\mathcal{D}}_{LCC}(F^s,M^s,v^s)),$$

$$\begin{gathered} {\mathcal{D}}_{\text{LCC}}(F^s,M^s,v^s)= \\ \frac{1}{P}\sum _{x\in \Omega }\frac{{\sum }_i{\left(\left(F_{x_i}^{\ast s}-\overline{F_x^{\ast s}}\right)\left(M_{x_i}^{\ast s}-\overline{M_x^{\ast s}}\right)\right)}^2}{\left({\sum }_i{\left(F_{x_i}^{\ast s}-\overline{F_x^{\ast s}}\right)}^2\right)\left({\sum }_i{\left(M_{x_i}^{\ast s}-\overline{M_x^{\ast s}}\right)}^2\right)+\tau }-1, \end{gathered}$$

$$\text{log }p_{\theta}(F|M)-\text{KL}\left[q_{\omega}(z|F,M)\|p(z|F,M)\right]=\\ \mathbb{E}_{z\sim q}\Big{[}\text{log }p_{\theta}(F|z,M)\Big{]}-\text{KL}\left[q_{\omega}(z|F,M)\|p(z)\right].$$

$$\text{KL}[\mathcal{N}(\mu_{\omega}(F,M),\Sigma_{\omega}(F,M))\|\mathcal{N}(0,I)]=\\ \frac{1}{2}\Big{(}\text{tr}(\Sigma_{\omega}(F,M))+\|\mu_{\omega}(F,M)\|-k-\text{log det}(\Sigma_{\omega}(F,M))\Big{)},$$

$${\color[rgb]{0,0,0}\mathbb{E}_{F,M\sim\cal{X}}\bigg{[}\mathbb{E}_{z\sim q}\Big{[}\text{log }p_{\theta}(F|z,M)\Big{]}-\text{KL}\left[q_{\omega}(z|F,M)\|p(z)\right]\bigg{]}.}$$

$$\mathbb{E}_{F,M\sim\cal{X}}\bigg{[}\mathbb{E}_{\epsilon\sim\mathcal{N}(0,I)}\Big{[}\text{log }p_{\theta}(F|z=\mu_{\omega}(F,M)+\\ \Sigma_{\omega}^{1/2}(F,M)\ast\epsilon,M)\Big{]}-\text{KL}\left[q_{\omega}(z|F,M)\|p(z)\right]\bigg{]}.$$

$$\mathbb{E}_{F,M\sim\cal{X}}\bigg{[}\mathbb{E}_{\epsilon\sim\mathcal{N}(0,I)}\Big{[}\sum_{s\in\mathcal{S}}\text{log }p_{\theta^{s}}(F^{s}|z=\mu_{\omega}(F,M)+\\ \Sigma_{\omega}^{1/2}(F,M)\ast\epsilon,M^{s})\Big{]}-\text{KL}\left[q_{\omega}(z|F,M)\|p(z)\right]\bigg{]}.$$

$${\hat{v}}^s=G_{{\sigma }_G}\ast v^s$$

$$\text{log}\hskip 2.84526ptp(v^{*})\propto\int_{\Omega}\sum_{i=1}^{\infty}{\color[rgb]{0,0,0}\frac{\sigma_{G}^{2i}}{2^{i}i!}}\left[\frac{\partial^{i}}{\partial\Omega^{i}}v^{*}\right]^{2}\hskip 2.84526ptd\Omega,$$

$$\hat{v}=\underset{v^{\ast }}{\mathrm{argmax}}\log p(v^{\ast }|v)=\underset{v^{\ast }}{\mathrm{argmax}}\log p(v|v^{\ast })+\log p(v^{\ast })$$

$$p_{\theta }(F|M)$$

$$dvd{v}^{\ast }dz$$

$$={\int }_{z}p(F|\hat{v},M)p(z)dz.$$

$$p_{\theta }(F|M)={\int }_z\prod _{s\in \mathcal{S}}p(F^s|{\hat{v}}^s,M^s)p(z)dz.$$

$$\operatorname*{arg\,min}_{\omega,\theta}\frac{1}{2}\Big{(}\text{tr}(\Sigma_{\omega})+\mu_{\omega}^{\top}\mu_{\omega}-k-\text{log det}(\Sigma_{\omega})\Big{)}\\ -\lambda\sum_{s\in\mathcal{S}}\mathcal{D}_{\textit{LCC}}(F^{s},M^{s},\hat{v}^{s}),$$

Full text

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.

I Introduction

Deformable image registration, the process of finding voxel correspondences in a pair of images, is an essential task in medical image analysis [1]. This mapping – the deformation field – can be used for example in pre-op / post-op studies, to find the same structures in images from different modalities or to evaluate the progression of a disease. The analysis of geometric changes in successive images is important for instance for diagnosing cardiovascular diseases and selecting the most suited therapies. A possible approach is to register sequential images and analyze the extracted deformations for example by parallel transport [2] or by creating an adapted low-dimensional subspace [3].

We propose a registration algorithm that learns a deformation model directly from training images. Inspired by recent generative latent variable models, our method learns a low-dimensional probabilistic deformation encoding in an unsupervised fashion. This latent variable space encodes similar deformations close to each other and allows the generation of synthetic deformations for a single image and the comparison and transport of deformations from one case to another.

I-A Deformable Image Registration

Traditionally, deformable registration is solved by numerical optimization of a similarity metric which measures the distance between the fixed and the deformed moving image. The moving image is warped given a predefined deformation model in order to get closer to the fixed image. Unfortunately, this results in an ill-posed problem which requires further regularization based on prior assumptions [1]. Various regularization energies have been proposed including elastic- [4, 5] or diffusion-based methods [6, 7, 8] (cf. [1]). Diffeomorphic transforms are folding-free and invertible. The enforcement of these properties in many medical applications has led to the wide use of diffeomorphic registration algorithms. Popular parametrizations of diffeomorphisms include the Large Deformation Diffeomorphic Metric Mapping (LDDMM) [9, 10, 11], a symmetric normalization approach [12] or stationary velocity fields (SVF) [13, 14].

In recent years, learning-based algorithms – notably Deep Learning (DL) – have been proposed to avoid long iterative optimization at test time. In general, one can classify these algorithms as supervised and unsupervised. Due to the difficulty of finding ground truth voxel mappings, supervised methods need to rely on predictions from existing algorithms [15, 16], simulations [17, 18, 19] or a combination of both [20, 21]. The latter can be achieved for example by projecting B-spline displacement estimations in the space of a statistical deformation model from which one can extract simulations by sampling of its components [20]. Diffeomorphic approaches predict patches of the initial momentum of LDDMMs [15] or dense SVFs [16]. Supervised methods are either limited by the performance of existing algorithms or the realism of simulations. Furthermore, retrieving deformations from existing algorithms on a large database is time-consuming and increases the training complexity.

Unsupervised approaches to registration aim to optimize an image similarity, often combined with a penalization or smoothing term (regularization). These learning approaches first appeared in the computer vision community [22, 23] and were recently applied to medical image registration [24, 25, 26, 27, 28]. Unlike traditional methods, learning-based approaches also can include task-specific information such as segmentation labels during training while not requiring those labels at test time. Instead of using an image similarity, Hu et al. [29] proposed to optimize the matching of labels based on a multi-scale DICE loss and a deformation regularization. Fan et al. [26] proposed to jointly optimize a supervised and unsupervised objective by regressing ground-truth deformation fields (from an existing algorithm), while simultaneously optimizing an intensity-based similarity criterion. The disadvantage of these semi-supervised approaches is that their training complexity is higher since label information needs to be collected, and for example deformations outside the segmented areas are not guaranteed to be captured. Most unsupervised approaches use B-spline grids or dense deformation fields, realized with spatial transformer layers (STN [30]) for an efficient and differentiable linear warping of the moving image. However, it has not been shown yet that these approaches lead to sufficiently regular and plausible deformations.

I-B Deformation Analysis and Transport

Understanding the deformation or motion of an organ goes beyond the registration of successive images. Therefore, it has been proposed to compare and characterize shape and motion patterns by normalizing deformations in a common reference frame [2, 31] and for example by applying statistical methods to study the variation of cardiac shapes [32]. In the diffeomorphic setting, various dimensionality reduction methods have been proposed. Vaillant et al. [33] modeled shape variability by applying PCA in the tangent space to an atlas image. Qiu et al. used a shape prior for surface matching [34]. While these methods are based on probabilistic inference, dimensionality reduction is done after the estimation of diffeomorphisms. Instead Zhang et al. [35] introduced a latent variable model for principle geodesic analysis that estimates a template and principle modes of variation while infering the latent dimensionality from the data. Instead of having a general deformation model capable of explaining the deformations of any image pair in the training data distribution, this registration approach still depends on the estimation of a smooth template. Using the SVF parametrization for cardiac motion analysis, Rohé et al. [3] proposed to build affine subspaces on a manifold of deformations, the barycentric subspaces, where each point on the manifold represents a 3-D image and the geodesic between two points describes the deformation.

For uncertainty quantification, Wassermann et al. [36] used a probabilistic LDDMM approach applying a stochastic differential equation and Wang et al. [37] employed a low-dimensional Fourier representation of the tangent space of diffeomorphisms with a normal assumption. While both approaches contain probabilistic deformation representations, they have not been used for sampling and the representations have not been learned from a large dataset.

In the framework of diffeomorphic registration, parallel transport is a promising normalization method for the comparison of deformations. Currently used parallel transport approaches are the Schild’s [38] or pole ladder [2, 39] using the SVF parametrization or approaches based on Jacobi fields using the LDDMM parametrization [40, 41]. In general, these approaches aim to convert and apply the temporal deformation of one subject to another subject. However, this transport process typically requires multiple registrations, including difficult registrations between different subjects.

I-C Learning-based Generative Latent Variable Models

Alternatively and inspired by recently introduced learning-based generative models, we propose to learn a latent variable model that captures deformation characteristics just by providing a large dataset of training images. In the computer vision community, such generative models as Generative adversarial networks (GAN) [42], stochastic variational autoencoders (VAE) [43] and adversarial autoencoders (AAE) [44] have demonstrated great performance in learning data distributions from large image training sets. The learned models can be used to generate new synthetic images, similar to the ones seen during training. In addition, probabilistic VAEs are latent variable models which are able to learn continuous latent variables with intractable posterior probability distributions (encoder). Efficient Bayesian inference can be used to deduce the posterior distribution by enforcing the latent variables to follow a predefined (simple) distribution. Finally, a decoder aims to reconstruct the data from that representation [43]. As an extension, conditional variational autoencoders (CVAE) constrain the VAE model on additional information such as labels. This leads to a latent variable space in which similar data points are mapped close to each other. CVAEs are for example used for semi-supervised classification tasks [45].

Generative models also showed promising results in medical imaging applications such as in classifying cardiac diseases [46] or predicting PET-derived myelin content maps from multi-modal MRI [47]. Recently, unsupervised adversarial training approaches have been proposed for image registration [21, 48, 28]. Dalca et al. [27] developed a framework which enforces a multivariate Gaussian distribution on each component of the velocity field for measuring uncertainty. However, these approaches do not learn global latent variable models which map similar deformations close to each other in a probabilistic subspace of deformations. To the best of the authors’ knowledge, generative approaches for registration which allow the sampling of new deformations based on a learned low-dimensional encoding have not been proposed yet.

I-D Probabilistic Registration using a Generative Model

We introduce a generative and probabilistic model for diffeomorphic image registration, inspired by generative latent variable models [43, 45]. In contrast to other probabilistic approaches such as [15, 27], we learn a low-dimensional global latent space in an encoder-decoder neural network where the deformation of a new image pair is mapped to and where similar deformations are close to each other. This latent space, learned in an unsupervised fashion, can be used to generate an infinite number of new deformations for any single image from the data distribution and not only for a unique template as in the Bayesian inference procedure for model parameter estimation in [35]. From this abstract representation of deformations, diffeomorphic deformations are reconstructed by decoding the latent code under the constraint of the moving image. To the best of the author’s knowledge, this method describes the first low-dimensional probabilistic latent variable model that can be used for deformation transport from one subject to another. Through applying a latent deformation code of one image pair on a new constraining image, deformation transport (and sampling from the latent space) is useful for instance for simulating cardiac pathologies or synthesizing a large number of pathological and healthy heart deformations for data augmentation purposes.

We use a variational inference method (a CVAE [45]) with the objective of reconstructing the fixed image by warping the moving image. The decoder of the CVAE is conditioned on the moving image to ease the encoding task: by making appearance information easily accessible in the decoder (in the form of the moving image), the latent space is more likely to encode deformation rather than appearance information. This implicit decoupling of deformation and appearance information allows to transport deformations from one case to another by pairing a latent code with a new conditioning image. The framework provides multi-scale estimations where velocities are extracted at each scale of the decoding network. We use the SVF parametrization and diffeomorphisms are extracted using a vector field exponentiation layer, based on the scaling and squaring algorithm proposed in [13]. This algorithm has been successfully applied in neural networks in our previous work [49] and in [27]. The framework contains a dense spatial transformer layer (STN) and can be trained end-to-end with a choice of similarity metrics: to avoid asymmetry, we use a symmetric and normalized local cross correlation criterion. In addition, we provide a generic formulation to include regularization terms to control the deformation appearance (if required), such as diffusion regularization in form of Gaussian smoothing [8]. During training, similarity loss terms for each scale and a loss term enforcing a prior assumption on the latent variable distribution are optimized by using the concept of deep supervision (cf. [50]). During testing, the low-dimensional latent encoding, multi-scale estimations of velocities, deformation fields and deformed moving image are retrieved in a single forward path of the neural network.

We evaluate our framework on the registration of cardiac MRIs between end-diastole (ED) and end-systole (ES) and provide an intensive analysis on the structure of the latent code and evaluate its application for transporting encoded deformations from one case to another.

This paper extends our preliminary work [49] by adding:

• •

Detailed derivations of the probabilistic registration framework including a generic regularization model.

• •

Deep supervision, multi-scale estimations and a normalized loss function to improve registration performances.

• •

Analysis of size and structure of the latent variable space.

• •

Evaluation of the deformation transport by comparing it to a state-of-the-art algorithm [2].

II Methods

In image registration, the goal is to find the spatial transformation $\mathcal{T}_{z}:\mathbb{R}^{3}\rightarrow\mathbb{R}^{3}$ which is parametrized by a $d$-dimensional vector $z\in\mathbb{R}^{d}$. The optimal values of $z$ are the ones which best warp the moving image $M$ in order to match the fixed image $F$ given the transformation $\mathcal{T}_{z}$. Both images $F$ and $M$ are defined in the spatial domain $\Omega\in\mathbb{R}^{3}$. Typically, $z$ is optimized by minimizing an objective function of the form: $\operatorname*{arg\,min}_{z}\mathcal{F}(z,M,F)=\mathcal{D}\left(F,M\circ\ \mathcal{T}_{z})+\mathcal{R}(\mathcal{T}_{z}\right)$, where $\mathcal{D}$ is a metric measuring the similarity between fixed $F$ and warped moving image $M\circ\ \mathcal{T}_{z}$. $\mathcal{R}$ is a spatial regularizer [1]. Recent unsupervised DL-based approaches [22, 24, 25] try to learn to maximize such a similarity metric $\mathcal{D}$ using stochastic gradient descent methods and a spatial transformer layer (STN [30]) for warping the moving image $M$.

In extension, we propose to model registration by learning a probabilistic deformation parametrization vector $z$ from a set of example image pairs $(M,F)$. Thereby, we constrain the low-dimensional $z$ to follow a prior distribution $p(z)$. In other words, our approach contains two key parts: a latent space encoding to model deformations and a decoding function that aims to reconstruct the fixed image $F$ from this encoded transformation – by warping the moving image $M$. In addition, this decoding function is generative as it allows to sample new deformations based on $p(z)$.

II-A Probabilistic model for multi-scale registration

We assume a generative probabilistic distribution for registration $p_{true}(F|M)$, capturing the deformation from $M$ towards $F$. We aim at learning a parameterized model $p_{\theta}(F|M)$ with parameters $\theta$ which allows us to sample new $F$’s that are similar to samples from the unknown distribution $p_{true}$. To estimate the posterior $p_{\theta}$ we use a latent variable model parameterized by $z$. Following the methodology of a VAE [43], we assume the prior $p(z)$ to be a multivariate unit Gaussian distribution with spherical covariance $I$:

[TABLE]

Using multivariate Gaussians offers a closed-form differentiable solution, however, $p(z)$ could take the form of other distributions. In this work, we parameterize deformation fields $\phi$ by stationary velocity fields (SVF), denoted by velocities $v$: $\phi=\text{exp}(v)$ [13]. These transformation maps $\phi$ are given as the sum of identity and displacements $u(x)$ for every position $x\in\Omega$: $\phi(x)=x+u(x)$. In the multi-scale approach, we define velocities $v^{s}$ at scale $s\in\mathcal{S}$ where $\mathcal{S}$ is the set of different image scales ($s=1$ describes the original scale for which we omit writing $s$ and $s=2,3,...$ the scale, down-sampled by a factor of $2^{s-1}$). For each scale $s$, a family of decoding functions $f_{v}^{s}$ is defined, parameterized by a fixed $\theta^{s}\subset\theta$ and dependent on $z$ and the moving image $M^{s}$:

[TABLE]

In the training, the goal is to optimize $\theta^{s}$ such that all velocities $v^{s}$ are likely to lead to warped moving images $M^{*s}$ that will be similar to $F^{s}$ in the training database. $M^{*s}$ is obtained by exponentiation of $v^{s}$ and warping of the moving image. Using Eq. 2, we can define the families of functions $f^{s}$:

[TABLE]

In order to express the dependency of $f^{s}$ on $z$ and $M^{s}$ explicitly, we can define a distribution $p(F^{s}|z,M^{s};\theta^{s})$. The product over the different scales gives us the output distribution:

[TABLE]

By using the law of total probability, this leads to the following stochastic process for computing $p_{\theta}(F|M)$ which is also visualized in Fig. 1a (cf. [45]):

[TABLE]

The likelihood $p_{\theta}(F|z,M)$ can be any distribution that is computable and continuous in $\theta$. In VAEs, the choice is often Gaussian, which is equivalent to adopting a sum-of-squared differences (SSD) criterion (cf. [43]). We propose instead to use a local cross-correlation (LCC) distribution due to its robustness properties and superior results in image registration compared to SSD (cf. [8, 51]). Thus, we use the following Boltzmann distribution as likelihood:

[TABLE]

where $v^{s}=f^{s}_{v}(z,M^{s};\theta^{s})$ are the velocities and $\lambda$ is a scalar hyperparameter. The symmetric $\mathcal{D}_{LCC}$ is defined as:

[TABLE]

with $P$ pixels $x\in\Omega$, the symmetrically warped images $M^{*s}=M^{s}\circ\text{exp}\left(v^{s}/2\right)$ and $F^{*s}=F^{s}\circ\text{exp}\left(-v^{s}/2\right)$. The bar $\overline{F_{x}}$ symbolizes the local mean grey levels of $F_{x}$ derived by mean filtering with kernel size $k$ at position $x$. $i$ is iterating through this $k\times k$-window. A small constant $\tau$ is added for numerical stability ($\tau=1e^{-15}$).

II-A1 Learning the constrained deformation encoding

In order to optimize the parameterized model over $\theta$ (Eq. 5), two problems must be solved: First, how to define the latent variables $z$, for example decide what information these variables represent. VAEs assume there is no simple interpretation of the dimensions of $z$ but instead assert that samples of $z$ are drawn from a simple distribution $p(z)$.

Second, the integral over $z$ is intractable since one would need to sample a too large number of $z$’s to get an accurate estimate of $p_{\theta}(F|M)$. Instead of sampling a large number of $z$’s, the key assumption behind VAEs is to sample only $z$’s that are likely to have produced $F$ and compute $p_{\theta}(F|M)$ only from those. To this end, one needs to compute the intractable posterior $p(z|F,M)$. Due to this intractability, in VAEs [43], the posterior is approximated by learning an encoding distribution $q_{\omega}(z|F,M)$, using a neural network with parameters $\omega$ (the encoder). This approximated distribution can be related to the true posterior using the Kullback-Leibler divergence (KL) which leads (after rearranging the terms) to the evidence lower bound (ELBO) of the log marginalized likelihood $\text{log }p_{\theta}(F|M)$ (cf. [43, 45]):

[TABLE]

The KL-divergence on the left hand side gets smaller the better $q_{\omega}(z|F,M)$ approximates $p(z|F,M)$ and ideally vanishes if $q_{\omega}$ is of enough capacity. Thus, maximizing $\text{log }p_{\theta}(F|M)$ is equivalent to maximizing the ELBO on the right hand side of the equation consisting of encoder $q_{\omega}$ and decoder $p_{\theta}$ which can be both optimized via stochastic gradient descent.

II-A2 Optimizing the ELBO

According to the right-hand side of Eq. 8, there are two terms to optimize, the KL-divergence of prior $p(z)$ and encoder distribution $q_{\omega}(z|F,M)$ and the expectation of the reconstruction term $\text{log }p_{\theta}(F|z,M)$. Since the prior is a multivariate Gaussian, the encoder distribution is defined as $q_{\omega}(z|F,M)=\mathcal{N}(z|\mu_{\omega}(F,M),\Sigma_{\omega}(F,M))$, where $\mu_{\omega}$ and $\Sigma_{\omega}$ are deterministic functions learned in an encoder neural network with parameters $\omega$. The KL-term can be computed in closed form as follows (constraining $\Sigma_{\omega}$ to be diagonal):

[TABLE]

where $k$ is the dimensionality of the distribution.

The expected log-likelihood $\mathbb{E}_{z\sim q}\left[\text{log }p_{\theta}(F|z,M)\right]$, the reconstruction term, could be estimated by using many samples of $z$. To save computations, we treat $p_{\theta}(F|z,M)$ as $\mathbb{E}_{z\sim q}\left[\text{log }p_{\theta}(F|z,M)\right]$ by only taking one sample of $z$. This can be justified as optimization is already done via stochastic gradient descent, where we sample many image pairs $(F,M)$ from the dataset $\mathcal{X}$ and thus witness different values for $z$. This can be formalized with the expectation over $F,M\sim\mathcal{X}$:

[TABLE]

To enable back-propagation through the sampling operation $q_{\omega}(z|F,M)$, the reparametrization trick [43] is used in practice, where $z=\mu_{\omega}+\epsilon\Sigma_{\omega}^{1/2}$ (with $\epsilon\sim\mathbf{N}(0,I)$). Thus, for image pairs $(F,M)$ from a training dataset $\cal{X}$ the actual objective becomes:

[TABLE]

After insertion of Eq. 4, the log of the product over the scales $s\in\mathcal{S}$ results in the sum of the log-likelihood distributions:

[TABLE]

II-B Introducing regularization on velocities

So far, we have considered that at each scale $s$, a velocity field $v^{s}$ is generated by a decoding function $f_{v}^{s}(z,M^{s};\theta^{s})$ through a neural network. To have a better control of its smoothness, we propose to regularize spatially $v^{s}$ through a Gaussian convolution with standard deviation $\sigma_{G}$:

[TABLE]

Gaussian smoothing was applied here, but it could be replaced by any quadratic Tikhonov regularizer or by any functional enforcing prior knowledge about the velocity field.

In the remainder, we show how the regularization of velocities can be inserted into the proposed probabilistic framework. To make the notation less cluttered, we drop the scale $s$ superscript in the velocity notations. Until now, the velocities $v$ have been handled as fixed parameters $v=f_{v}(z,M;\theta)$. We can equivalently assume that velocities $v$ are random variables with a Dirac posterior probability : $p_{\theta}(v|z,M)\sim\delta_{f_{v}(z,M;\theta)}(v)$.

We now introduce the random variable $v^{*}$ which represents the regularized velocities as shown in Fig. 1b. This quantity is linked to the regular velocities $v$ through a Gaussian distribution $p(v|v^{*})=G(v^{*},1)$ such that $v$ is close to $v^{*}$ in terms of $L_{2}$ norm. Furthermore, we define a diffusion-like regularization prior on $v^{*}$ [52]:

[TABLE]

taking into account the Taylor expansion of the Fourier transform of the Gaussian. The maximum a posteriori of the regularized velocities $\hat{v}$ is then obtained through Bayes law:

[TABLE]

which in this case is equivalent to solving the Heat equation [52] and leads to a Gaussian convolution: $\hat{v}=G_{\sigma_{G}}\ast v$.

Finally, we conveniently assume that the posterior probability of $v^{*}$ is infinitely peaked around its mode, i.e. p(v^{*}|{\color[rgb]{0,0,0}v})\sim\delta_{\hat{v}}(v^{*}) (assumption sometimes made for the Expectation-Maximization algorithm [53]). In the decoding process, we can now marginalize out the velocity variables $v$ and $v^{*}$ by integrating over both such that only $\hat{v}$ remains:

[TABLE]

Thus, the proposed graphical model leads to a decoder working with the regularized velocity field $\hat{v}$ instead of the $v$ generated by the neural network. When combining regularized velocities $\hat{v}^{s}$ at all scales, we get:

[TABLE]

This can be optimized as before and leads to Gaussian convolutions at each scale if considering diffusion-like regularization. Thus, the multi-scale loss function per training image pair ($F$,$M$) for one sample $\epsilon$ is defined as (cf.  Eq. 10):

[TABLE]

where $\hat{v}^{s}$ depends on $v^{s}$ and therefore on $\theta$ (cf. Eq. 2 and 11).

II-C Network architecture

The encoder-decoder neural network takes the moving and the fixed image as input and outputs the latent code $z$, velocities $v$, the deformation field $\phi$ and the warped moving image $M^{*}$. The last three are returned at the different scales $s$. The encoder consists of strided convolutions while the bottleneck layers ($\mu$, $\sigma$, $z$) are fully-connected. The deconvolution layers in the decoder were conditioned by concatenating each layer’s output with sub-sampled versions of $\mathbf{M}$. Making appearance information of the moving image easily accessible for the decoder, allows the network to focus on deformation information – the differences between moving and fixed image – that need to pass through the latent bottleneck. While it is not guaranteed that the latent representation contains any appearance information, it comes at a cost to use the small bottleneck for appearance information. At each decoding scale, a convolutional layer reduces the number of filter maps to three. Then, a Gaussian smoothing layer (cf. Eq. 11) with variance $\sigma_{G}^{2}$ is applied on these filter maps. The resulting velocities $v^{s}$ (a SVF) are exponentiated by the scaling and squaring layer [49] in order to retrieve the diffeomorphism $\phi^{s}$ which is used by a dense STN to retrieve the warped image $M^{*s}$. The latent code $z$ is computed according to the reparametrization trick. During training, the network parameters are updated through back-propagation of the gradients with respect to the objective Eq. 10, defined at each multi-scale output. Finally during testing, registration is done in a single forward path where $z$ is set to $\mu$ since we want to execute registration deterministically. One can also think of drawing several $z$ using $\sigma$ and use the different outputs for uncertainty estimation as in [27] which we do not further pursue in this work. The network architecture can be seen in Fig. 2a. Besides registration, the trained probabilistic framework can be also used for the sampling of deformations as shown in Fig. 2b.

III Experiments

We evaluate our framework on cardiac intra-subject registration. End-diastole (ED) frames are registered to end-systole (ES) frames from cine-MRI of healthy and pathological subjects. These images show large deformations. Additionally, we evaluate the learned encoding of deformations by visualizing the latent space and transporting encoded deformations from one patient to another. All experiments are in 3-D.

III-1 Data

We used the 334 ED-ES frame pairs of short-axis cine-MRI sequences. 184 cases were acquired from different hospitals and 150 cases were used from the Automatic Cardiac Diagnosis Challenge (ACDC) at STACOM 2017 [54], mixing congenital heart diseases with images from adults. We used 234 cases for training and for testing 100 cases from ACDC, that contain segmentation and disease class information. The testing set contained 20 cases of each of the following cardiac diseases: dilated cardiomyopathy (DCM), hypertrophic cardiomyopathy (HCM), previous myocardical infarction (MINF), abnormal right ventricle (RV) and healthy (Normal). All images were resampled with a spacing of $1.5\times 1.5\times 3.15$ mm and cropped to a size of $128\times 128\times 32$ voxels, by equally removing voxels from all sides. These dimensions were chosen to save computation time and are not a limitation of the framework.

III-2 Implementation details

Our neural network consisted of four encoding convolutional layers with strides (2, 2, 2, 1) and three decoding deconvolutional layers. Each scale contained two convolutional layers and a convolutional Gaussian layer with $\sigma_{G}=3mm$ (kernel size 15) in front of an exponentiation and a spatial transformer layer using trilinear interpolation (cf. Fig. 2a). The dimensionality of the latent code $z$ was set to $d=32$ as a compromise of registration quality and generalizability (cf. experiment on latent vector dimensionality). The number of trainable parameters in the network was $\sim$420k. LeakyReLu activation functions and L2 weight decay of $1*10^{-4}$ were applied on all layers except the last convolutional layer in each scale where a tanh activation function was used in order to avoid extreme velocity values during training. All scales were trained together, using linearly down-sampled versions of the input images for the coarser scales. In all experiments, the number of iterations in the exponentiation layer was set to $N=4$ (evaluated on a few training samples according to the formula in [13]). During the training, the mean filter size of the LCC criterion was $k=9$. The loss hyper parameter was empirically chosen as $\lambda=5000$ such that the similarity loss was optimized while the latent codes roughly had zero means and variances of one. We applied a learning rate of $1.5*10^{-4}$ with the Adam optimizer and a batch size of one. For augmentation purposes, training image were randomly shifted, rotated, scaled and mirrored. The framework has been implemented in Tensorflow using Keras111https://keras.io/. Training took $\sim$24 hours and testing a single registration case took 0.32s on a NVIDIA GTX TITAN X GPU.

III-3 Registration

We compare our approach with the LCC-demons (Dem, [8]) and the ANTs software package using Symmetric Normalization (SyN, [12]) with manually tuned parameters (on a few training images) and the diffeomorphic DL-based method VoxelMorph [27] (VM) which has been trained using the same augmentation techniques as our algorithm. For the latter, we set $\sigma=0.05$, $\lambda=50000$ and applied a reduced learning rate of $5*10^{-5}$ for stability reasons while using more training epochs. Higher values for $\lambda$ led to worse registration accuracy. We also show the improvement of using a multi-scale approach (with 3 scales, S3) compared to a single-scale objective (S1).

We measure registration performance with the following surrogates: intensity root mean square error (RMSE), DICE score, 95%-tile Hausdorff distance (HD in mm). To quantify deformation regularity, we show the determinant of the Jacobian qualitatively, while we also computed the mean magnitude of the gradients of the determinant of the Jacobian (Grad Det-Jac). We decided to report this second-order description of deformations to better quantify differences in smoothness among the different methods, which are not obvious by taking the mean of the determinant of the Jacobian as bigger and smaller values tend to cancel each other out. DICE and HD scores were evaluated on the following anatomical structures: myocardium (LV-Myo) and epicardium (LV) of the left ventricle, left bloodpool (LV-BP), right ventricle (RV) and LV+RV (Fig. 4).

Table I shows the mean results and standard deviations of all algorithms. In terms of DICE scores, our algorithm using three scales (Our S3) shows the best performances on this dataset while the single-scale version (Our S1) performed similarly compared to the LCC-demons and the SyN algorithm. Hausdorff distances were significantly improved using both of our algorithms. Detailed registration results are shown in Fig. 3. Interestingly, we found that the SyN algorithm showed marginally better DICE scores than the LCC-demons which has been also reported on brain data [8].

Qualitative registration results of a pathological (HCM) and a healthy case (Normal) are presented in Fig. 5a222Qualitative registration results for all five diseases are also presented in Fig. 10-12 (available in the supplementary files /multimedia tab).. The warped moving image (with and wihout grid overlay) and the determinant of the Jacobian (Det. Jac.) are shown. Displacements are visualized using the color encoding as typical for the optical flow in computer vision tasks. Middle and coarse scale outputs of our multi-scale method are shown in Fig. 5b. We computed the determinant of the Jacobian using SimpleITK333http://www.simpleitk.org/ and found that for all methods no negative values were observed on our test dataset. Compared to the other algorithms, our approach produced smoother and more regular deformations as qualitatively shown by the determinant of the Jacobian in Fig. 5a and quantitatively by the significantly smaller mean gradients of the determinant of the Jacobian (Table I)444Visualization of the gradients of the Jacobian determinant are presented in Fig. 13 (available in the supplementary files /multimedia tab).. Despite the fact of being diffeomorphic, the voxelmorph algorithm produced more irregular deformation fields compared to all other algorithms. Our single-scale approach resulted in slightly smoother deformations which is probably due to the fact that it performed less accurately in compensating large deformations.

We applied the Wilcoxon signed-rank test with $p<0.001$ to evaluate statistical significance of the differences in the results of Fig. 3. This method is chosen as a paired test without the assumption of normal distributions. For all metrics, the results of our multi-scale algorithm (Our S3) showed significant differences compared to the results of all other methods (including Our S1). With respect to our single-scale algorithm (Our S1), only the differences in DICE scores were not statistically significant in comparison with the LCC-demons (Dem).

Note, that higher DICE and HD scores can be achieved by choosing a higher latent dimensionality (cf. Experiment III-4), which however comes at the cost of a more complex encoding space, making analysis tasks more difficult. We also tested the first version of voxelmorph [25] on our dataset. We chose to show the results of the latest version [27] due to the fact that this version is diffeomorphic and that its DICE and HD results were better (cf. [49]).

III-4 Deformation encoding

For evaluating the learned latent space, we investigated (a) the effects of the size of the latent vector on the registration accuracy, (b) the structure of the encoded space by visualizing the distribution of cardiac diseases and showing simulated deformations along the two main axes of variations and (c) we applied our framework on deformation transport and compare its performance with a state-of-the-art algorithm.

Latent Vector Size

In Fig. 6 we analyzed the influence of the size of the latent code vector with respect to registration accuracy in terms of DICE and HD scores. With a relatively small latent size of $d=8$, competitive accuracy is achieved. With an increasing dimensionality, performance increases but reaches a plateau eventually. This behavior is expected, since CVAEs tend to ignore components if the dimensionality of the latent space is too high [45]. For the cardiac use case, we chose $d=32$ components as a trade-off between accuracy and latent variable size.

Disease Distribution and Generative Latent Space

In this experiment, we used disease information and encoded z-codes of the test images to visualize the learned latent space. Using linear CCA (canonical correlation analysis), we projected the $z$-codes (32-D) to a 2-D space by using the two most discriminative CCA components. Fig. 7 shows that the 100 test sets are clustered by classes in this space. Taking the 8 most discriminative CCA components into account, the classification accuracy of the five classes is 83% with 10-fold cross-validation using support vector machines (SVM). In a second experiment, we applied principal component analysis (PCA) on the $z$-codes of the training dataset. We simulated deformations by sampling equally distributed values in the range of $\pm 2.5$ standard deviations of the two largest principal components and extracting the $z$-codes through inverse projections. Fig. 8 shows reconstructed displacements and deformed images when applying these generated $z$-codes on a random test image. One can see the different influences of the two eigenvalues. The first eigenvalue (horizontal) focuses on large deformations while the second one focuses on smaller deformations as the right ventricle. The results of these two experiments which are solely based on applying simple linear transformations, suggest that deformations that are mapped close to each other in the deformation latent space have similar characteristics.

Deformation Transport

Pathological deformations can be transported to healthy subjects by deforming the healthy ES frames using pathological ED-ES deformations. Our framework allows for deformation transport by first registering the ED-ES frames of a given pathological case (Step 1), which we call prediction in this experiment. Secondly, we use the $z$-codes from these predictions along with the ED frame of a healthy subject to transport the deformations (Step 2, cf. Fig. 2b). Note, that this procedure does not require any inter-subject registrations.

We compare our approach with the pole-ladder algorithm (PL [2]). All intra- and inter-subject registrations required by the pole ladder were performed using the LCC-demons [8]. For the inter-subject pairs, we aligned the test data with respect to the center of mass of the provided segmentation and rotated the images manually for rigid alignment. This alignment step was done only for the pole ladder experiment555The pipeline for the parallel transport experiment using the pole ladder algorithm is presented in Fig. 14 (available in the supplementary files /multimedia tab)..

Qualitative results are shown in Fig. 9 where the predicted deformations of one hypertrophy (A, HCM) and one cardiomyopathy (B, DCM) case (step 1) were transported to two healthy (Normal) subjects (step 2, targets C and D). Note that our algorithm automatically determines orientation and location of the heart. In Table II, we evaluated the average ejection fraction (EF) of the ED-ES deformation prediction of the pathologies (step 1) and the average EF after transport to normal subjects (step 2). Hereby, we assume that EFs, as a relative measure, stay similar after successful transport (such that the absolute difference, EF step 1 - EF step 2, is small).

The table shows the average of transporting 5 HCM and 5 DCM cases to 20 normal cases (200 transports). For our algorithm, the absolute differences in EFs are much smaller for DCM cases and similarly close in HCM cases in comparison to the pole ladder. All test subjects were not used during training. The EF is computed based on the segmentation masks (warped with the resulting deformation fields). Besides, it can be seen, that predictions done by the demons are underestimating the EFs for HCM cases which should be $>$40% according to the ACDC data set specifications.

IV Discussion and Conclusions

We presented an unsupervised multi-scale deformable registration approach that learns a low-dimensional probabilistic deformation model. Our method not only allows the accurate registration of two images but also the analysis of deformations. The framework is generative, as it is able to simulate deformations given only one image. Furthermore, it provides a novel way of deformation transport from one subject to another by using its probabilistic encoding. In the latent space, similar deformations are close to each other. The method enables the addition of a regularization term which leads to arbitrarily smooth deformations that are diffeomorphic by using an exponentiation layer for stationary velocity fields. The multi-scale approach, providing velocities, deformation fields and warped images in different scales, leads to improved registration results and a more controlled training procedure compared to a single-scale approach.

We evaluated the approach on end-diastole to end-systole cardiac cine-MRI registration and compared registration performance in terms of RMSE, DICE and Hausdorff distances to two popular algorithms [8, 12] and a learning-based method [27], which are all diffeomorphic. While the performance of our single-scale approach was comparable to the LCC-demons and the SyN algorithm, our multi-scale approach (using 32 latent dimensions) showed statistically significant improvements in terms of registration accuracy. Generally, our approach produced more regular deformation fields, which are significantly smoother than the DL-based algorithm. Using our method with a non-generative U-net style network [55] without a deformation encoding performed similarly compared to the proposed generative model. Adding supervised information such as segmentation masks in the training procedure as in [29, 26] led to a marginal increase in terms of registration performance ($\sim$1-2% in DICE scores), so we decided that the performance gain is not large enough in order to justify the higher training complexity. Theoretically, our method allows measurement of registration uncertainty as proposed in [27] which we did not further investigate in this work.

The analysis of the deformation encoding showed that the latent space projects similar deformations close to each other such that diseases can be clustered. Disease classification could be potentially enforced in a supervised way as in [46]. Furthermore, our method showed comparable quantitative and qualitative results in transporting deformations with respect to a state-of-the-art algorithm which requires the difficult step of inter-subject registration that our algorithm does not need.

It is arguable if the simple assumption of a multivariate Gaussian is the right choice for the prior of the latent space (Eq. 1). Possible other assumptions such as a mixture of Gaussians are subject to future work. The authors think that the promising results of the learned probabilistic deformation model could be also applicable for other tasks such as evaluating disease progression in longitudinal studies or detecting abnormalities in subject-to-template registration. An open question is how the optimal size of the latent vector changes in different applications. In future work, we plan to further explore generative models for learning probabilistic deformation models.

Acknowledgments

Data used in preparation of this article were obtained from the EU FP7-funded project MD-Paedigree and the ACDC STACOM challenge 2017 [54]. This work was supported by the grant AAP Santé 06 2017-260 DGA-DSH, and by the Inria Sophia Antipolis - Méditerranée, ”NEF” computation cluster. The authors would like to thank Xavier Pennec for the insightful discussions and Adrian Dalca for the help with the Voxelmorph [27] experiments.

**Disclaimer: ** This feature is based on research, and is not commercially available. Due to regulatory reasons its future availability cannot be guaranteed.