Bi-level Probabilistic Feature Learning for Deformable Image Registration

1 Introduction

Image registration [Maintz and Viergever, 1998], transforming different sets of images into one common coordinate system, is one of the most fundamental tasks in computer vision. Especially, deformable image registration which builds a dense correspondence between image pairs has wide applications in medical image analysis such as multi-modality fusion, anatomical change diagnosis, and population modeling. The high degrees of freedom for the solution space (deformation maps) and great variations on source/target image pairs are major challenges for this issue.

Conventional image registration approaches typically formulate the discrepancy between image pairs as well as the prior constraints on images and deformations into an energy function whose optimization finds the solution to registration [Ashburner, 2007]. These approaches are able to generate plausible deformation maps with desired mathematical properties, e.g., invertibility, and topology-preserving, even when significant appearance variations exist. Unfortunately, the iterative optimization process demands calculating gradients over the high dimensional deformation and image/feature spaces at each step, resulting in extremely high computational expenses.

Recently, researchers bring deep learning networks that gain great success in classification/segmentation to image registration in order to achieve efficient prediction for deformations. Balakrishnan et al. develop the VoxelMorph (VM) network similar to UNet structure [Balakrishnan et al., 2019], and the works of [Dalca et al., 2019] and [Shen et al., 2019] combine the learning with the diffeomorphic constraint for topology-preserving deformations. Hu et al. [Hu et al., 2019] develop a multi-scale network to compute the registration field from convolutional feature pyramids. These works learn deep features in the image domain and then apply the learned parameters as the one-step prediction for deformation. They just directly learn parameters based on predesigned training loss/regularization to output the deformation filed. So it is hard to adaptively enforce registration information for the front-end feature learning phase. Consequently, images deviating from the training greatly affect the performance, yielding unstable registration.

In this work, deformation learning embraces feature learning in a bi-level optimization framework. We formulate the energy for deformation in the feature domain as one level of optimization embedded with the other level of maximum a posterior (MAP) optimization for features. Hence, two challenging issues in learning-based approaches, i.e.,, learning optimal networks for deformation and learning optimal features for registration, are cooperatively resolved in one model. We detail our contributions as:

• We establish a bi-level model to simultaneously address deformable registration and feature learning. The upperlevel of optimization learns the deformation field in the feature domain while the lower-level learns the features required for image registration. We alternately solve the deformation filed based on the upper-level subproblem and adjust the features based on the lower-level subprob-

lem, giving robust and efficient registration.

• We pose the feature learning in the lower-level as probabilistic MAP estimation with Gaussian priors. This level of optimization adaptively enforces task-based constraints for the feature extraction phase, producing optimal features for the upper-level optimization.
• We develop an unsupervised end-to-end training scheme for the deep networks respective to the two levels of optimization. The scheme, training the networks using the loss function on the feature domain and image domain respectively, runs more efficiently than training a complicated image-to-deformation network.

Extensive experiments on 3D brain MRI registration tasks demonstrate that our approach achieves state-of-the-art performance in terms of accuracy, robustness, and efficiency.

2 RelatedWork

Recently, feature learning has been employed in many image analysis tasks. Taking advantage of the network to learn the high-level features has shown impressive results [Liu et al., 2018a] [Liu et al., 2019a]. Research in [Banerjee et al., 2019] uses a conditional encoder-decoder network for the segmentation task. The work of [Sun et al., 2018] estimates the optical flow on the multi-scale feature space. Replacing handcrafted features with more discriminative learnable feature space has been extensively developed [Cheng et al., 2018] [Liu et al., 2019b] [Liu et al., 2018b]. Researchers are tending towards superseding the feature engineering or original image with learned convolution filters, within end-to-end trainable architectures. The feature learning designed for the registration task has also been explored in [Hu et al., 2019] [de Vos et al., 2019] [Liu et al., 2020]. However, these networks are designed to learn deterministic features.

Traditional deformable registration techniques include the b-splines model with control points [Sun et al., 2014], elastictype models [Bajcsy and Kovacic, 1989] and diffeomorphic registration. Diffeomorphic frameworks [Ashburner, 2007] use smooth flow fields to represent the deformation, regularization is typically introduced as part of the ordinary differential equation constraining on the vector fields. Due to a large number of parameters, numerical optimization becomes computationally costly, which is the main limitation of these traditional methods. Recently, learning-based methods [Dosovitskiy et al., 2015] have been widely applied. The research in [Yang et al., 2017] proposes a supervised learning approach to rapidly predict 3D deformable registrations, achieving an order of magnitude speed-up. But ground truth registration fields are hard and expensive to obtain. Inspired by the VoxelMorph [Balakrishnan et al., 2019], researchers have focused on replacing costly numerical optimization with global function optimization over the training data in an unsupervised way. The research in [Hu et al., 2019] employs an unsupervised coarse-to-fine dual-stream registration network, enabling the capability for handling significant deformations. Moreover, some researches[Shen et al., 2019] propose to estimate the velocity fields or momentum fields, which can be used to obtain diffeomorphic transformations.

3 The ProposedMethod

In this section, we first introduce a novel bi-level framework to incorporate feature optimizing strategy into deformable image registration. Bi-level optimization [Colson et al., 2005] is a special optimization scheme, in which one optimization task is nested within another. The outer level is given by the constraint inner level problem, referred as upperlevel tasks and lower-level tasks respectively. Then, we propose our deep architectures, the probabilistic feature learning module and the deformable registration module to solve the two nested tasks.

3.1 Bi-level Registration Framework

The proposed bi-level model jointly optimizes the features and registration fields. In this novel model, the upper-level problem denotes the deformable registration, while the lowerlevel problem is a maximum posterior probability problem. The outer deformable registration problem is constrained by the inner feature optimization. Specifically, given a source image Is and a target image It with a spatial domain Ω ∈ Rd , we aim at minimizing:

where fs, ft are the feature representations of the source and target image, ϕ : Ω×R → Ω is the final deformation field. ED is a data matching term, forcing the image similarity. ER is a regularization term, encouraging the smoothness of the deformation fields. p(fs|Is, ft|It,ϕ) denotes the posterior probability of features based on the observation input source, target image and the deformation field ϕ. Note that, to better solve the deformable registration problem, we employ the stationary velocity fields. Such that the final registration field ϕ is obtained via integration of a series of velocity fields over time. We employ regularization into the registration fields by introducing the integration of the velocity fields.

The proposed bi-level model considers the feature optimization and deformable registration model simultaneously, which is different from other models in essence. Our formulation is much more robust than the typical registration model since it takes feature optimization into consideration.

3.2 Probabilistic Feature LearningModule

To learn the most suitable feature representations that are invariant to noise and uninformative intensity-variations, we propose to learn the feature via MAP estimation, corresponding to the lower-level problem in Eq. (1). Specially, we use the posterior probability to learn the most meaningful nonlinear mapping from the input intensities to feature representation. We aim to estimate the most suitable feature for the registration task as follows:

where ϕ is the corresponding registration field, and I is the input source or target image. The first term is the data likelihood term and the second one is the prior term.

We assume the prior p(f) to be a multivariate unit normal distribution with covariance I. However, the posterior probability is difficult to compute, we use an approximation posterior probability qθ(f ) parametrized by θ to speed it up. We try to minimize the Kullback-Leibler (KL) divergence between the really posterior distribution and approximation distribution, equivalent to maximizing the Evidence Lower Bound (ELBO) [Kingma and Welling, 2014] as follows:

where the first term means samples from the approximate distribution to perform the registration tasks. The second term is the KL term, also called Gauss prior term, which encourages the two probability distribution to be close. We model the approximate posterior qθ(f ) as a multivariate normal:

where N(.; µ, Σ) is the multivariate normal distribution with the mean µ and covariance Σ. Then we optimize the above evidence lower bound.

We optimize the variational lower bound by learning the network parameters of the probabilistic feature learning networks using backpropagation methods. The feature learning networks take the input images/features and output the approximate posterior probability parameters of feature representations, representing mean µf and covariance Σf , then sample a new feature, which is processed to generate the deformation fields at different scale. Specifically, to generate feature representations/distribution at the kth scale, we use layers of convolutional filters to downsample the features at the previous k + 1th pyramid level, by a factor of 2. We employ three-layers feature learning network, which includes convolutional layers with 16 filters or 32 filters. Each convolutional layer uses leaky ReLU function and 3×3×3 kernel.

3.3 Deformable RegistrationModule

Diffeomorphic deformable registration [Ashburner, 2007] can capture large deformations and provide many desirable mathematical properties, such as invertibility, globally oneto-one smooth, and topology-preserving. In this work, we develop an optimization scheme to solve the diffeomorphic deformable registration, corresponding to the upper-level problem in Eq. (1), as follows:

where each deformation field ϕ is obtained via the integration of current V, governed by the ordinary differential equation φt(x, t) = V(φ(x, t), t) over time t, with identity map φ(0) = Id. And φ1 is the sought-for registration field, with k standing for scale.

To generate the velocity field sequence, we employ an efficient deep residual architecture, named by the alignment network, to solve the Eq. (6) as follows:

where Wk are the learnable parameters of the alignment network, and the deformation field ϕ is defined through the ordinary differential equation of V, corresponding to regularization module. We implement the regularization module using the efficient scaling and squaring [Dalca et al., 2019] method. As for the alignment network, it processes the features and the deformation field from the previous scale to generate the current deformation field. The alignment networks at different blocks share the same network structure, including three convolutional layers with {48, 32, 16} filters respectively. All convolutional layers are followed by a leaky ReLU activation function except the one that outputs the registration field. And we use the 3×3×3 kernel for all the convolutional layers.

image-to-image [2] registration. Standard deviations are in bracket.

4 End-to-end Training Strategy

All the modules are differentiable, thus we learn the network parameters with the end-to-end training. The proposed propagation network is jointly trained using stochastic gradient descent-based methods in an unsupervised way. The loss functions are as follows:

Loss on feature space. The training loss on feature space consists of the task-specific reconstruction loss and the KL loss, forcing the encoded distribution to be close to the prior probability distribution. The hypothesis is that to better optimize the reconstruction loss, the feature learning networks at each scale make use of the provided posterior probability information of features, resulting in the feature optimizing. For the task-specific reconstruction loss, we scale the image pairs and warp [Jaderberg et al., 2015] the down-sampled images with the deformation fields at different scales, then compute the similarity loss using the Local Normalized Cross Correlation. Therefore, the deformation fields at each scale can influence the previous featuring learning. Specifically, we use different window sizes to compute the local normalized correlation coefficient at each scale, with a smaller window size for the lower resolution. From the zeroth to the second scale, the window sizes are set to {5, 5, 7}, respectively. For the KL term, we obtain Σf and µf from the feature learning module, then compute it in the closed-form as:

with Σf to be diagonal, and k is const. We use the multi-scale training loss, computing the sum of the losses at different scales.

Loss on final deformation field. We also consider the information on the image domain, using the final deformation field as follows:

we take the Local Normalized Cross Correlation as image similarity loss, with the window size as 9. For the regularization term, we also employ a diffusion regularizer on the final deformation field.

5 Experiments

In this section, we first introduce our experiment setup, including datasets, evaluation metrics, baseline methods, and our implementation. Next, to demonstrate the superiority of our method, we compare it with the state-of-the-art deformable registration techniques on the accuracy, efficiency as well as diffeomorphism preservation of the deformation fields. Then we explore the impact of the inner task of our paradigm.

Data Preparation. Evaluations are conducted in two different sights, one aligning all the source data to a common atlas, called image-to-atlas registration, and the other addressing general registration between two arbitrary volumes, called images-to-image registration. For both of these two cases, 368 T1 weighted MR volumes from three publicly available datasets: ADNI [Mueller et al., 2005], ABIDE [Di Martino et al., 2014] and OASIS [Marcus et al., 2007] are selected and split into 281, 17, and 70 for training, validation, and testing, respectively. Also, the unseen PPMI [Marek et al., 2011] dataset containing 59 scans is employed for testing. Specifically, for the image-to-atlas registration, we use the atlas provided by [Balakrishnan et al., 2019] as target, and for image-to-image case, the target images are randomly selected from datasets. Considering the large disparity among different datasets, standard pre-processing operations were conducted, including motion correction, NU intensity correction, normalization, skull stripping, and affine registration, as described in the research [Balakrishnan et al., 2019]. To reduce memory usage, the images are cropped to 160 × 192 × 224.

Evaluation Metrics. To achieve a more comprehensive evaluation, both the average Dice score [Dice, 1945] over registered testing pairs and the Jacobian matrix over the computed deformation are considered as evaluation metrics, to evaluate the anatomical overlap correspondences of the registered volume pairs and the smoothness of the deformation fields. To calculate the Dice score, segmentation is performed with FreeSurfer on each of the testing volumes to extract 30 anatomical structures, on which the average Dice is calculated. The Jacobian matrix Jφ(x) =∇φ(x) captures the local properties of φ around voxel x, such as stretching and rotating, negative determinants mean the loss of the one-to-one mapping [Ashburner, 2007]. So the local deformation is diffeomorphic at the locations where |x : Jφ(x) > 0|. We count all the folds, the negative Jacobian locations, and use Folds to represent the number of folds.

BaselineMethods. We compare our method with state-ofthe-art registration techniques, including three optimizationbased tools: Elastix [Klein et al., 2010], Symmetric Normalization (SyN) [Avants et al., 2008], NiftyReg [Sun et al.,

2014], and two learning-based methods: VoxelMorph [Balakrishnan et al., 2019] and its diffeomorphic variant [Dalca et al., 2019] (referred as VM and VM-diff, respectively). For the SyN algorithm, we use the version implemented in the ANTs [Avants et al., 2011] package. The parameter settings of the conventional methods are as follows. For Elastix, we run B-spline registration with Mattes Mutual Information as cost function and set the control point spacing to 16 voxels. Four scales are used with 500 iterations per scale. For the SyN algorithm, we take Cross Correlation as the similarity measure metric and use the SyN step size of 0.25, Gaussian parameters (9, 0.2), at three scales with 201 iterations each. As for NiftyReg, we use the Normalized Mutual Information cost function. We run it with 12 threads using 1500 iterations. We run Elastix, ANTs (SyN), and NiftyReg on a PC with i7-8700 (@3.20GHz, 32G RAM), while learning-based methods on NVIDIA TITAN XP.

Implementation. We implement our networks with TensorFlow [Abadi et al., 2016] package. It takes about 12 hours to train our model from scratch with 28100 iterations. During training, we use Adam optimizer [Kingma and Ba, 2015] with a learning rate of 1e-4 . We set the batch size as 1. Computation on the full resolution may easily exhaust the memory, thus we choose to output a half-resolution smooth enough deformation field and up-sample it to full-resolution.

5.1 Comparison Results

First, we quantitatively evaluate the accuracy, rationality and time consumption of all these techniques for both cases of image-to-atlas and image-to-image registration. Tab. 1 depicts the stability of the methods in terms of the Dice score on the different datasets, where higher values and lower variance indicate a more accurate and stable registration. Our method gives an obvious lower variance with a comparable mean of Dice for both these two cases, showing stronger stability. As shown in Tab. 2, only VM-diff and our method can decrease the number of folds to zero, preserving the diffeomorphism. Besides, as Tab. 3 shows, our approach requires less running time, benefiting of the well-designed network architectures. Dealing with half-resolution rather than the original scale further accelerates the registration process.

To have a better understanding of the alignment results, we illustrate the Dice score of 30 anatomical structures in Fig. 2. Limited by space, besides our method, we only present ANTs (SyN), VM and VM-diff as the representatives for the optimization-base and learning-based techniques. We can see that compared with the conventional method ANTs (SyN), the deep methods VM, VM-diff give evenly accuracy but perform much less stable among different anatomical segmentations. While our deep model achieves a good balance between accuracy and stability. In summary, the robustness, accuracy, and speed are all the key performance indexes for the registration method. And extensive validation experiments indicate that our method performs favorably against state-of-the-art methods in practice.

Fig. 3 shows our representative registration results, from which we can see that our approach can ideally handle large changes in shapes. As shown, our method can ideally preserve the contour of anatomical structures, guaranteeing the topology of the registered volumes.

5.2 Ablation Studies

We investigate the roles of probabilistic feature learning, the lower-level optimization in our bi-level model, and the scales of learned features. Experiments were conducted on the testing data from multiple datasets the mixed datasets (referred to as Test) and the unseen PPMI dataset (referred to as PPMI). We substitute our probabilistic learning feature pyramids (PFL) with handcrafted image pyramids as the case of 'w/o PFL' and compare the performance gap between these two cases. Tab. 4 lists the registration accuracy in terms of both the Dice score and Local Normalized Correlation Coefficient (LNCC) at 2, 3 and 4 scales of features. We can see that our probabilistic feature learning evidently increases the accuracy in terms of both metrics. As the table shows, the features with three scales output the superior accuracy over the other scales. we observe in our experiments that more scales have more risks of overfitting when a limited number of training examples are available. Therefore, we adopt the probabilistic feature learning with three scales in all our experiments.

Fig. 4 shows the confidence maps of learned features for one slice of the target and source images, and the deformation fields in order to give a more intuitive illustration of the probabilistic features and uncertainty estimation, we show uncertainty maps. The features with higher confidences are coincident with prominent image structures. Therefore, the registration upon aligning these structures is able to produce stabler outputs than deterministic learning.

6 Conclusions

We introduce a bi-level optimization model to simultaneously address the deformable registration and feature optimization. It takes the energy for deformation computation as the upperlevel optimization while formulates the maximum a posterior for features as the lower-level optimization. Such deformation learning policy with adaptive registration feature extraction is completely different from existing straightforward learning-based image registration methods. Then we employ efficient deep architectures to simultaneously propagate deformation fields and perform feature optimization. The losses on the feature space and image domain are utilized in our unsupervised end-to-end learning framework. We conduct two groups of image registration experiments on 3D brain MRI datasets including image-to-atlas and image-to-image registrations. Extensive results show that our method achieves state-of-the-art performance with extreme efficiency. We have demonstrated our performance on the mon-modal registration task, and future validation remains on multi-modal registration.

Acknowledgments

This work is partially supported by the National Natural Science Foundation of China (Nos. 61922019, 61672125, 61733002 and 61772105), and LiaoNing Revitalization Talents Program (XLYC1807088). We thank Adrian V. Dalca for sharing code to simulate the image for the Boxplots indicating Dice scores for anatomical structures.