Deep Group-wise Variational
Diffeomorphic Image Registration

Image registration, the process of aligning images, is a fundamental problem in medical image analysis [32]. For example, it can be used to align image volumes in longitudinal studies [25], or to align images acquired temporally during contrast enhancement [27].

A common way to register images is to apply a transformation, such as an affine transformation, B-splines [28], or a deformation (vector) field [29], and to optimize its parameters using some form of numerical optimization. The objective is to minimize misalignment defined by a loss, such as a squared error, negative normalized cross correlation, or more advanced dissimilarity measures such as negative normalized mutual information [24]. Additionally, unrealistic deformations are mitigated by regularization or by using a constraint family of deformations.

Methods that constrain deformation fields to diffeomorphisms, such as LDDMM [4], Diffeomorphic Demons [31], and SyN [1], guarantee preservation of topology, and have shown to be very effective.

Recently, it has been shown that deep learning models can be trained to perform image registration, leveraging on the increased availability of medical imaging data [34, 2]. Once a model is trained, deformation fields can be obtained by inference, which is much faster than aligning images through an iterative optimization process. In [13] and [14], diffeomorphic neural network-based methods were proposed, further bridging the gap between learning-based and classical approaches. Although deep learning methods show great potential, they are restricted to image pairs or require a heuristically constructed template [8].

In this work, we extend learning-based methods to enable fast group-wise image registration. We build upon the variational framework proposed by [14] and stack multiple images along the channel axis of the input and output of the network. In addition, we propose a likelihood based on normalized mutual information and a prior which yields explicit control over the viscous fluid energy. We quantitatively evaluate our model and compare it with Elastix [18] and Voxelmorph [3].

Depending on the task, we define images that require registration (i.e. moving images) as observations of $\mathbf{x}$ and the target images (i.e. fixed images) as observations of $\mathbf{y}$. We are not restricted to a predefined fixed timepoint and can describe all of our data as observations of $\mathbf{x}$, effectively letting $\mathbf{y}$ be an unobserved latent variable. We will refer to this case as all-moving group-wise image registration, and the case where one timepoint is chosen as fixed image as all-to-one group-wise image registration. Observe that we can obtain the former from the latter by constraining the deformation $\phi_{\mathbf{z}_{p}}^{-1}=\text{Id}$ at a fixed timepoint $p$ to be the identity. Moreover, we can view pair-wise registration as a particular $K=2$ instance in our group-wise framework, where one timepoint is kept fixed and the other one can move.

We let a 3D U-net [26, 2, 3], output the parameters $\boldsymbol{\mu}_{\mathbf{z}}$ and $\mathbf{\Sigma}_{\mathbf{z}}$ of the approximate posterior $q_{\theta}(\mathbf{z}|\mathbf{x},\mathbf{y})$, as illustrated in Figure 2. Fixed and moving image volumes are concatenated along the channel dimension and form the input of the network. Figure 3 gives a more detailed description of the network architecture, including filter sizes and channel counts. All layers are followed by a Leaky ReLU [20] with 0.2 slope, except for the $\boldsymbol{\mu}_{\mathbf{z}}$ and $\mathbf{\Sigma}_{\mathbf{z}}$ heads which are respectively followed by a hyperbolic tangent scaled with $\alpha=15$ that limits the output to $[-\alpha,\alpha]$ voxels and a Softplus ¹¹1$\text{Softplus}(x)=\log(\exp(x)+1)$ activation to ensure positive variances.

We train for 100k iterations using Adam [16] with a batch size of 1 and a learning rate of $1e^{-5}$ ($\beta_{1}=0.9,\beta_{2}=0.999$) decayed with cosine annealing [19]. Inputs are normalized using 1-99 percentile normalization [23] and training samples consist of $96\times 96\times 96$ patches for all timepoints at a randomly sampled location. For the hyperparameters we choose $\lambda_{u}=1$ and $\lambda_{v}=0.01$.

Using the proposed GroupMorph model, we have performed intra-patient registration of multiple 3D scans acquired at different timepoints in two different ways. First, we performed all-to-one registration where the scan acquired at the first time point was used as fixed image and all other scans of the same patient as moving images. Thereafter, we performed all-moving registration, where we compute deformation fields for all images and the register to a geodesic average.

To compare performance of the proposed GroupMorph model with state-of-the-art conventional and neural-based methods, we performed registration using Elastix [18] and VoxelMorph [3, 14] . For registration with Elastix we use configuration described for registration of DCE-MRI descriped in [15]. For fair comparison with VoxelMorph, we use the same loss, network architecture and hyper-parameters in VoxelMorph and GroupMorph. We only vary the channels in the first and last layers to match required dimensions and set $\lambda_{v}=0$ for VoxelMorph.

Experiments were performed using both breast MRI and chest CT scans. For breast MR scans no landmarks were available, hence evaluation was performed with root mean square error (RMSE) and NMI. For chest CT, a set of 75 landmarks was provided at all timepoints within DIR-Lab, hence performance was evaluated by measuring the average distance between the landmarks after registration and over all timepoint combinations. Moreover, for both sets we report percentages of negative Jacobian determinants ($\%\ |J_{\phi}|<0$) that indicate undesirable folding. Additionally, we evaluate the speed of registrations. Results are listed in Table 1 ²²2We measured time to calculate all required deformation fields for Elastix on two 10-core Intel Xeon Silver 4114 2.20GHz CPUs and of the neural-based methods on a NVIDIA GeForce RTX 2080 GPU.. They show that the group-wise models perform better in terms of RMSE and NMI on both sets, especially in the all-moving case. We find that diffeomorphic registration models (indicated with ‘-diff’) lead to lower amount of negative Jacobian determinants, which means that less folding occurs and indicates that the fields are smoother than those of non-diffeomorphic counterpart. Average registration times per series (Table 1) show that registration times of the GroupMorph model are significantly lower than the registration times of the baselines, especially compared with Elastix.

Figure 4 illustrates the axial slices of the first volume in the DIR-Lab dataset. Subtraction images illustrate improved alignment for all registration models, especially at borders and in particular for all-moving group-wise models. From the Jacobian determinant fields, we can observe particularly strong local expansions and compression in the all-moving models and most notably in the non-diffeomorphic variant.

In this study, we proposed a variational and diffeomorphic learning-based registration method for application in group-wise image registration. We evaluated our approach using intra-patient registration with sets of multiple 3D scans of two different modalities (MR and CT) showing different anatomies (breast and chest). In addition, we provided a likelihood based on normalized mutual information, a well-performing image similarity metric in registration, between multiple images and a prior that allows for explicit control over the viscous fluid energy to effectively control regularization of deformation fields.

In spite of using diffeomorphic transformations, we observed some negative Jacobian determinants likely resulting from approximation errors. Hence, future work could investigate options for further improvements. It can be expected that training with more data, further tuning of the network architecture and hyper-parameters will lead to higher performance. Other potential advancements include a 4D network [9] to more explicitly model the temporal component, fine tuning of $\mathbf{z}$ after inference (see dotted grey arrow in Figure 2), time-dependent velocity fields to allow more flexible diffeomorphic transformations, and a bi-directional cost [12] and non-diagonal covariance [14] to further improve smoothness of the deformation fields.

To conclude, we found that our GroupMorph model can simultaneously register multiple scans with performance similar to conventional registration with Elastix and learning-based VoxelMorph, enabling close to real-time group-wise registration without delaying clinical reading.

Authors would like to thank Jorrit Glastra, Tim Henke and Rob Hesselink for insightful discussions and support.