Given time-dependent velocity field $\tilde{v}_{t}\in\tilde{V}$, the diffeomorphism $\phi^{-1}_{t}\in\widetilde{{\rm Diff}}(\Omega)$ in the finite-dimensional Fourier domain can be computed as

where $\tilde{\mathcal{D}}\tilde{\phi}^{-1}_{t}$ is a tensor product $\tilde{\mathcal{D}}\otimes\tilde{\phi}^{-1}_{t}$, representing the Fourier frequencies of a Jacobian matrix $\tilde{\mathcal{D}}$. The $\ast$ is a circular convolution with zero padding to avoid aliasing. We truncate the output of the convolution in each dimension to a suitable finite set to avoid the domain growing to infinity.

Geodesic shooting algorithm states that the transformation can be uniquely determined by integrating a partial differential equation with a given initial velocity $v_{0}$ forward in time. In contrast to the original LDDMM that optimizes over a collection of time-dependent velocity fields, geodesic shooting estimates an optimal initial velocity $v_{0}$. In this work, we adopt an efficient variant of geodesic shooting defined in Fourier spaces. We first discretize the Jacobian and divergence operators using finite difference approximations (particularly the central difference scheme) and then compute their Fourier coefficients. Detailed derivations can be found in the appendix section of FLASH (Zhang and Fletcher, 2019). The Fourier representation of the geodesic constraint that satisfies Euler-Poincaré differential (EPDiff) equation is (Zhang and Fletcher, 2015)

where $\tilde{\mathcal{L}}$ is a symmetric, positive-definite differential operator that is a function of parameter $\alpha$ (details are in Sec. 3.1). Here $\tilde{\mathcal{K}}$ is the inverse operator of $\tilde{\mathcal{L}}$, and $\star$ is the truncated matrix-vector field auto-correlation ¹¹1The output signal maintains bandlimited after the auto-correlation operates on zero-padded input signal followed by truncating it back to the bandlimited space.. The ${\rm ad}^{\dagger}$ is an adjoint operator to the negative Jacobi–Lie bracket of vector fields, ${\rm ad}_{\tilde{v}}\tilde{w}=-[\tilde{v},\tilde{w}]=\tilde{\mathcal{D}}\tilde{v}\ast\tilde{w}-\tilde{\mathcal{D}}\tilde{w}\ast\tilde{v}$. The operator $\tilde{\nabla}\cdot$ is the discrete divergence (computed by summing the Fourier coefficients of different directions over in $\tilde{\mathcal{D}}$) of a vector field.

Consider a source image $S$ and a target image $T$ as square-integrable functions defined on a torus domain $\Omega=\mathbb{R}^{d}/\mathbb{Z}^{d}$ ($S(x),T(x):\Omega\rightarrow\mathbb{R}$). The problem of diffeomorphic image registration is to find the geodesic (shortest path) of diffeomorphic transformations $\phi_{t}\in{\rm Diff}(\Omega):\Omega\rightarrow\Omega,t\in[0,1]$, such that the deformed image $S\circ\phi^{-1}_{1}$ at time point $t=1$ is similar to $T$.

The objective function can be formulated as a dissimilarity term plus a regularization that enforces the smoothness of transformations

The $\text{Dist}(\cdot,\cdot)$ is a distance function that measures the dissimilarity between images. Commonly used distance metrics include sum-of-squared difference ($L_{2}$ norm) of image intensities (Beg et al., 2005), normalized cross-correlation (NCC) (Avants et al., 2008), and mutual information (MI) (Wells III et al., 1996). The $\phi_{1}^{-1}$ denotes the inverse of deformation that warps the source image $S$ in the spatial domain when $t=1$. Here $\gamma$ is a weight parameter balancing between the distance function and regularity term. The distance term stays in the full spatial domain, while the regularity term is computed in bandlimited space.

In contrast to previous approaches, our proposed model is parameterized in a bandlimited velocity space $\tilde{V}$, with parameter $\alpha$ enforcing the smoothness of transformations.

Assuming an independent and identically distributed (i.i.d.) Gaussian noise on image intensities, we obtain the likelihood

where $\sigma^{2}$ is the noise variance and $M$ is the number of image voxels. The deformation $\phi_{1}^{-1}$ corresponds to $\tilde{\phi}_{1}^{-1}$ in Fourier space via the Fourier transform $\mathcal{F}(\phi_{1}^{-1})=\tilde{\mathcal{\phi}}_{1}^{-1}$, or its inverse $\phi_{1}^{-1}=\mathcal{F}^{-1}(\tilde{\mathcal{\phi}}_{1}^{-1})$. The likelihood is defined by the residual error between a target image and a deformed source at time point $t=1$. We assume the definition of this distribution is after the fact that the transformation field is observed through geodesic shooting; hence is not dependent on the regularization parameters.

Analogous to (Wang et al., 2018), we define a prior on the initial velocity field $\tilde{v}_{0}$ as a complex multivariate Gaussian distribution, i.e.,

where $|\cdot|$ is matrix determinant. The Fourier coefficients of $\tilde{\mathcal{L}}$ is, i.e., $\tilde{\mathcal{L}}=\left(\alpha\tilde{\mathcal{A}}+1\right)^{3}$, $\tilde{\mathcal{A}}(\xi_{1},\ldots,\xi_{d})=-2\sum_{j=1}^{d}\left(\cos(2\pi\xi_{j})-1\right)$. Here $\tilde{\mathcal{A}}$ denotes a negative discrete Fourier Laplacian operator with a $d$-dimensional frequency vector $(\xi_{1},\ldots,\xi_{d})$, where $d$ is the dimension of the bandlimited Fourier space.

Combining the likelihood in Eq. (4) and prior in Eq. (5) together, we obtain the negative log posterior distribution on the deformation parameter parameterized by $\tilde{v}_{0}$ as

Next, we optimize Eq. (6) over the regularization parameter $\alpha$ and the registration parameter $\tilde{v}_{0}$ by maximum a posterior (MAP) estimation using gradient descent algorithm. Other optimization schemes, such as BFGS (Polzin et al., 2016), or the Gauss-Newton method (Ashburner and Friston, 2011) can also be applied.

Gradient of $\alpha$. To simplify the notation, first we define $f(\tilde{v}_{0})\triangleq-\ln\,p(\tilde{v}_{0}\,|\,S,T,\sigma^{2},\alpha)$. Since the discrete Laplacian operator $\tilde{\mathcal{L}}$ is a diagonal matrix in Fourier space, its determinant can be computed as $\displaystyle\prod_{j=1}^{d}(\alpha\tilde{\mathcal{A}}_{j}+1)^{3}$. Therefore, the log determinant of $\tilde{\mathcal{L}}$ operator is

We then derive the gradient term $\nabla_{\alpha}f(\tilde{v}_{0})$ as

Gradient of $\tilde{v}_{0}$. We compute the gradient with respect to $\tilde{v}_{0}$ by using a forward-backward sweep approach developed in (Zhang et al., 2017). Steps for obtaining the gradient $\nabla_{\tilde{v}_{0}}f(\tilde{v}_{0})$ are as follows:

1. (i)

   Forward integrating the geodesic shooting equation Eq.(2) to compute $\tilde{v}_{1}$,

2. (ii)

   Compute the gradient of the energy function Eq. (3) with respect to $\tilde{v}_{1}$ at $t=1$,

   $\displaystyle\nabla_{\tilde{v}_{1}}f(\tilde{v}_{0})=\tilde{\mathcal{L}}^{-1}(\alpha)\left(\frac{1}{\sigma^{2}}(S\circ\phi_{1}^{-1}-T)\cdot\nabla(S\circ\phi_{1}^{-1})\right).$

   (8)

3. (iii)

   Bring the gradient $\nabla_{\tilde{v}_{1}}f(\tilde{v}_{0})$ back to $t=0$ by integrating adjoint Jacobi fields backward in time (Zhang et al., 2017),

   $$\frac{d{\hat{v}}}{dt}=-{\rm ad}^{\dagger}_{\tilde{v}}{\hat{h}},\quad\frac{d{\hat{h}}}{dt}=-\hat{v}-{\rm ad}_{\tilde{v}}\hat{h}+{\rm ad}^{\dagger}_{\hat{h}}\tilde{v},$$

   (9)

   where $\hat{v}\in V$ are introduced adjoint variables with an initial condition $\hat{h}=0,\hat{v}=\nabla_{\tilde{v}_{1}}f(\tilde{v}_{0})$ at $t=1$.

A summary of the optimization is in Alg. 1. It is worthy to mention that our low-dimensional Bayesian framework developed in bandlimited space dramatically speeds up the computational time of generating training regularity parameters by approximately ten times comparing with high-dimensional frameworks.

Now we are ready to introduce our network architecture by using the estimated $\alpha$ and given image pairs as input training data. Fig. 2 shows an overview flowchart of our proposed learning model. With the optimal registration regularization parameter $\alpha^{opt}$ obtained from an image pair (as described in Sec. 3.1), a two-stream CNN-based regression network takes source images, target images as training data to produce a predictive regularization parameter. We optimize the network $\mathcal{H}(S,T;W)$ with the followed objective function,

where $\text{Reg}(W)$ denotes the regularization on convolution kernel weights $W$. $\text{Err}[\cdot,\cdot]$ denotes the data fitting term between the ground truth and the network output. In our model, we use $L_{2}$ norm for both terms. Other network architectures, e.g. 3D Residual Networks (3D-ResNet) (Hara et al., 2017) and Very Deep Convolutional Networks (3D-VGGNet) (Simonyan and Zisserman, 2014; Yang et al., 2018) can be easily applied as well.

In our network, we input the 3D source and target images into separate channels that include four convolutional blocks. Each 3D convolutional block is composed of a 3D convolutional layer, a batch normalization (BN) layer with activation functions (PReLU or ReLU), and a 3D max-pooling layer. Specifically, we apply $5\times 5\times 5$ convolutional kernel and $2\times 2\times 2$ max-pooling layer to encode a batch (size as $B$) of source and target images ($128^{3}$) to feature maps ($16^{3}$). After extracting the deep features from source and target channels, we combine them into a fusion channel, which includes three convolutional blocks, a fully connected layer, and an average pooling layer to produce the network output.

We run experiments on both 2D synthetic dataset and 3D real brain MRI scans.

For both 2D and 3D datasets, we split the images by using $70\%$ as training images, $15\%$ as validation images, and $15\%$ as testing images such that no subjects are shared across the training, validation, and testing stage. We evaluate the hyperparameters of models and generate preliminary experiments on the validation dataset. The testing set is only used for computing the final results.

Fig. 3 displays our MAP estimation of registration results including appropriate regularity parameters on 2D synthetic images. The middle panel of Fig. 3 reports the convergence of $\alpha$ estimation vs. ground truth. It indicates that our low-dimensional Bayesian model provides trustworthy regularization parameters that are fairly close to ground truth for network training. The bottom panel of Fig. 3 shows the convergence graph of the total energy for our MAP approach.

Fig. 4 further investigates the consistency of our network prediction. The left panel shows estimates of regularization parameter at multiple scales, i.e., $\alpha=0.1,1.0,10.0$, over $900$ 2D synthetic image pairs respectively. The right panel shows the mean error of image differences between deformed source images by transformations with predicted $\alpha$ and target images. While there are small variations on estimated regularization parameters, the registration results are very close (with averaged error at the level of $10^{-5}$).

Fig. 5 shows examples of 2D pairwise image registration with regularization estimated by MAP and our predictive deep learning model. We obtain the regularization parameter $\alpha=11.34$ (MAP) vs. $\alpha=13.20$ (network prediction), and $\alpha=5.44$ (MAP) vs. $\alpha=6.70$ (network prediction). The error map of deformed images indicates that both estimations obtain fairly close registration results.

Fig. 6 displays deformed 3D brain images overlaid with transformation grids for both methods. The parameter estimated from our model produces a comparable registration result. From our observation, a critical pattern between the optimal $\alpha^{opt}$ and its associate image pairs is that the value of $\alpha$ is relatively smaller when large deformation occurs. This is because the image matching term (encoded in the likelihood) requires a higher weight to deform the source image.

Fig. 7 investigates the consistency of our network prediction over three different datasets of 3D brain MRI. The left panel shows the absolute value of numerical differences between predicted regularization parameters and MAP estimations. The right panel shows the voxel-wise mean error of image differences between deformed images by transformations with predicted $\alpha$ and deformed images by MAP. While slight numerical difference on estimated regularization parameters exists, the 3D deformed images are fairly close (with averaged voxel-wise error at the level of $10^{-6}$).

Fig. 8 visualizes three views of the deformed brain images (overlay with propagated segmentation label) that are registered by MAP and our prediction. Our method produces a registration solution, which is highly close to the one estimated by MAP. The propagated segmentation label fairly aligns with the target delineation for each anatomical structure. While we show different views of 2D slices of brains, all computations are carried out fully in 3D.

Fig. 9 reports the volume overlapping of nine anatomical structures for both methods, including Cor(cortex), Puta (putamen), Cere (cerebellum), Caud (caudate), gyrus, Stem (brain stem), Precun (precuneus), Cun (cuneus), and Hippo (hippocampus). Our method produces comparable dice scores comparing with MAP estimations. This indicates that the segmentation-based registration by using our estimation achieves comparable registration performance with little to no loss of accuracy.

Table. 1 quantitatively reports the averaged time and memory consumption of MAP estimation in full spatial image domain and our method. The proposed predictive model provides appropriate regularization parameters approximately $1000$ times faster than the conventional optimization-based registration method with a much lower memory footprint.