The likelihood $p\left({\cal D}\mid w\right)$ specifies the relationship between the data and the transformation parameters by means of a similarity metric. In probabilistic image registration, it usually takes the form of a Boltzmann distribution (Ashburner, 2007):

where $\mathcal{E}_{\text{data}}$ is the similarity metric and $\mathcal{H}$ an optional set of hyperparameters.

Local cross-correlation (LCC), which is invariant to linear intensity scaling, is a popular similarity metric but not meaningful in a probabilistic context. For this reason, instead of the sum of the voxel-wise product of intensities, like in standard LCC, we opt for the sum of voxel-wise squared differences of images standardised to zero mean and unit variance inside a local neighbourhood of five voxels. This way, we can benefit from robustness under linear intensity transformations, as well as desirable properties of a Gaussian mixture model (GMM)of intensity residuals, i.e. robustness to outlier values caused by acquisition artefacts and misalignment over the course of registration (Le Folgoc et al., 2017).

Let $\overline{F}$ and $\overline{M(w)}$ be respectively the fixed and the warped moving image with intensities standardised to zero mean and unit variance inside a neighbourhood of five voxels. For each voxel, the intensity residual $r_{i}=\overline{F}-\overline{M(w)}$, $i\in\{1,\dots,N^{3}\}$, is assigned to the $l$-th component of the mixture, $1\leq l\leq L$, if the categorical variable $c_{i}\in\{1,\dots,L\}$ is equal to $l$, in which case it follows a normal distribution $\mathcal{N}\left(0,\beta_{l}^{-1}\right)$¹¹1In order to reduce the notation clutter we omitted the voxel index for the fixed and moving images.. The component assignment $c_{i}$ follows a categorical distribution and takes value $l$ with probability $\varrho_{l}$. We use the same GMMof intensity residuals on a global basis rather than per neighbourhood. In all experiments it has $L=4$ components, which we determine to be sufficient for a good model fit.

We also use the scalar virtual decimation factor $\alpha$ to account for the fact that voxel-wise residuals are not independent. This prevents over-emphasis on the data term and allows to better calibrate uncertainty estimates (Groves et al., 2011; Simpson et al., 2012). The full expression of the image similarity term is given by:

In Bayesian models, the transformation parameters are typically regularised with use of a multivariate normal prior that ensures smoothness:

where $\lambda_{\text{reg}}$ is a scalar parameter that controls the regularisation strength, and $\mathrm{L}$ is the matrix of a differential operator. Here we assume that $\mathrm{L}$ represents the gradient operator, which penalises the magnitude of the 1 derivative of a velocity field. Note that $\left(\mathrm{L}w\right)^{\intercal}\mathrm{L}w=\|\mathrm{L}w\|^{2}\coloneqq$.

The regularisation weight $\lambda_{\text{reg}}$ can either be fixed or estimated from data. The latter has been done successfully only for transformation parametrisations with a relatively low number of degrees of freedom, e.g. B-splines (Simpson et al., 2012) and a sparse parametrisation (Le Folgoc et al., 2017). In case of an SVF, where the number of degrees of freedom is orders of magnitude higher, the problem is more difficult. However, a reliable method to adjust regularisation strength based on data is crucial, as both the output transformation and registration uncertainty are highly sensitive to regularisation. In order to infer the regularisation strength, we specify a log-normal prior on the scalar regularisation energy $\chi^{2}\sim\text{Lognormal}\left(\mu_{\chi^{2}},\sigma_{\chi^{2}}^{2}\right)$, and derive a prior on the underlying SVF:

where $\nu=3N^{3}$ is the number of degrees of freedom, i.e. the count of transformation parameters in all three directions. Given (semi-)informative hyperpriors on $\mu_{\chi^{2}}$ and $\sigma^{2}_{\chi^{2}}$, which we discuss in the next section, we can adjust the regularisation strength to the input images. The full expression of the regularisation term is given by:

It is worth noting that the traditional $L_{2}$ regularisation with a fixed regularisation weight in Equation 5 actually belongs to this family of regularisation losses. If we specify a gamma prior instead of a log-normal prior on the scalar regularisation energy $\chi^{2}\sim\Gamma\left(\nicefrac{{\nu}}{{2}},\nicefrac{{\lambda_{\text{reg}}}$, we get:

We set the likelihood GMMhyperpriors similarly to Le Folgoc et al. (2017), with the mixture precision parameters $\beta=\left(\beta_{1},\dots,\beta_{L}\right)$ assigned independent log-normal priors $\beta_{l}\sim\text{Lognormal}\left(\mu_{\beta_{l}},\sigma^{2}_{\beta_{l}}\right)$ and the mixture proportions $\varrho=\left(\varrho_{1},\dots,\varrho_{L}\right)$ with an uninformative Dirichlet prior $\varrho\sim\text{Dir}\left(\kappa\right)$, where $\kappa=\left(\kappa_{1},\dots,\kappa_{L}\right)$.

Regularisation parameters require informative priors due to the difficulty of learning the regularisation strength based on a single image pair. Because of a gamma prior on the regularisation energy $\exp\left(\mu_{\chi^{2}}\right)\sim\Gamma\left(\nicefrac{{\nu}}{{2}},\nicefrac$, we can rely on the familiar regularisation weight $\lambda_{\text{init}}$ to initialise the logarithm of the regularisation energy $\mu_{\chi^{2}}$ to the expected value of the logarithm of the gamma distribution, i.e. $\mathbb{E}\left[\mu_{\chi^{2}}\right]=\psi\left(\nicefrac{{\nu}}{{2}}\right)-$, where $\psi$ is the digamma function. The value of this expression is sharply peaked if the number of degrees of freedom $\nu$ is large, which yields a very informative prior on $\mu_{\chi^{2}}$. More details on how to calculate the expected value of the logarithm of a gamma distribution can be found in Section A.

The choice of a hyperprior on the scale parameter $\sigma_{\chi^{2}}$, which controls the amount of deviation of $\log\chi^{2}$ from the location parameter $\mu_{\chi^{2}}$, is more intuitive. Here we use a log-normal prior $\sigma_{\chi^{2}}^{2}\sim\text{Lognormal}\left(\eta,\varsigma^{2}\right)$.

The model is implemented in PyTorch. For all experiments we use three-dimensional T2-FLAIR MRIbrain scans and subcortical structure segmentations from the UK Biobank dataset (Sudlow et al., 2015). Input images are pre-registered with the affine component of drop2 (Glocker et al., 2008)²²2https://github.com/biomedia-mira/drop2 and resampled to $N=128$ isotropic voxels of length $1.82\text{\,}\mathrm{mm}$ along every dimension.

We use $2^{12}$ steps to integrate SVFs. In order to start optimisation with small displacements, $\mu_{w}$ is initialised to zero, which corresponds to an identity transformation, $\sigma_{w}$ to half a voxel length in each direction and $u_{w}$ to a tenth of a voxel length in each direction. We are mainly interested in the approximate variational posterior in order to initialise the SG-MCMCalgorithm, so the rank parameter is set to $R=1$. We use the Adam optimiser with a step size of $1\times 10^{-2}$ for the approximate variational posterior parameters $\mu_{w}$, $\log\sigma_{w}^{2}$, and $u_{w}$. Training is run until the ELBOvalue stops to increase, which requires approximately 1,024 iterations.

In the likelihood model, we use $\kappa=0.5$ for an uninformative Jeffreys prior on the mixture proportions, while the mixture precision hyperparameters are set to $\mu_{\beta_{l}}=0.0$ and $\sigma_{\beta_{l}}=2.3$. The model is much less sensitive to the value of the likelihood hyperparameters than the regularisation hyperparameters, which are calibrated to guarantee diffeomorphic transformations sampled from the approximate variational posterior $q\left(w\right)\sim\mathcal{N}\left(\mu_{w},\Sigma_{w}\right)$. The local transformation is diffeomorphic only in locations where the Jacobian determinant is positive (Ashburner, 2007), so we aim to keep the number of voxels where the Jacobian determinant is non-positive $|\det J_{\varphi^{-1}}|\leq 0$ close to zero. We calibrate the location hyperparameter $\lambda_{\text{init}}$ in every experiment, while the scale hyperparameters are set to $\eta=2.8$ and $\varsigma=5.0$.

SAEMconvergence is known to be conditional on decreasing step sizes (Delyon et al., 1999). For this reason, we use a small step size decay of $1\times 10^{-3}$ and the Adam optimiser with a step size of $2\times 10^{-1}$ for the GMM hyperparameters $\log\beta^{-0.5}$ and $\log\varrho$, and $1\times 10^{-2}$ for the regularisation hyperparameters $\mu_{\chi^{2}}$ and $\log\sigma_{\chi^{2}}$. In case of parameters whose value is constrained to be positive, we state the step size used on the logarithms. In practice, we did not observe the result to be dependent on these step sizes.

First we evaluate the proposed regularisation. We compare it to a gamma prior on $\lambda$, i.e. $\lambda\sim\Gamma\left(s,r\right)$, where $s$ and $r$ are the shape and the rate parameters respectively, set to uninformative values $s=r=\nicefrac{{\nu}}{{2}}$ (Simpson et al., 2012).

We compare the output of VIwhen using fixed regularisation weights $\lambda_{\text{reg}}\in\{0.1,1.2\}$, the baseline method for learnable regularisation strength, and our regularisation loss. The result on a sample pair of input images is shown in Figure 1. For the baseline method, the learnt regularisation strength is too high, which effectively prevents the alignment of images. This indicates that previous schemes for inference of regularisation strength from data are inadequate when the transformation parametrisation involves a very large number of degrees of freedom. In case of $\lambda_{\text{reg}}=0.1$, the resulting transformation is not diffeomorphic. The output when using our regularisation loss with $\lambda_{\text{init}}=1.2$ strikes a balance between the baseline and $\lambda_{\text{reg}}=0.1$, where there is an overemphasis on the data term.

In Figure 2, we show the output of VIfor two pairs of images which require different regularisation strengths for accurate alignment. We choose a fixed image and two moving images $M_{1}$ and $M_{2}$, with one visibly different and the other similar to the fixed image. We also analyse the result when using fixed regularisation weights $\lambda_{\text{reg}}=0.2$, which leads to non-diffeomorphic transformations, and $\lambda_{\text{reg}}=2.0$, which produces smooth transformations but, in case of $M_{1}$, at the expense of accuracy. The proposed regularisation, initialised with $\lambda_{\text{init}}=2.0$, helps to prevent oversmoothing.

We run a number of experiments to evaluate the uncertainty estimates and better understand the differences between uncertainty output by various non-rigid registration methods in practice:

1. 1.

   We compare the uncertainty estimates output by VI, SG-MCMC, and VoxelMorph on inter-subject brain MRIdata from UK Biobank qualitatively and quantitatively, by calculating the Pearson correlation coefficient between the displacement and label uncertainties;

2. 2.

   We compare the uncertainty estimates when the SG-MCMCalgorithm is initialised to different transformations;

3. 3.

   We compare the result when using non-parametric SVFsand SVFsbased on B-splines to parametrise the transformation.

In order to make sampling from SG-MCMCefficient, we determine the largest step size that guarantees diffeomorphic transformations as defined in Section 6.1 and set it to $\tau=4\times 10^{-1}$. A single Markov chain requires only $4\text{\,}\mathrm{GB}$ of memory, so two Markov chains are run in parallel, each initialised to a different sample from the approximate variational posterior. We discard the first 100,000 samples from each chain to allow the Markov chains to reach the stationary distribution.

To draw samples from the true posterior of transformation parameters remains a difficult problem even with a large number of simplifying model assumptions. If the true posterior were approximately Gaussian, VIwould provide a good approximation thereof and the lower uncertainty output by SG-MCMCwould indicate that the output samples are autocorrelated even with subsampling. However, if the true posterior is not Gaussian, then the posterior output by VIis ill-fitting, and the samples output by SG-MCMCcover multiple modes near the mode that corresponds to VI, which makes a good case for the use of SGLD.

In practice, the quality of uncertainty estimates is also sensitive to the validity of model assumptions. These include coinciding image intensities up to the expected spatial noise offsets and ignoring spatial correlations between residuals. The first assumption is valid in case of mono-modal registration but the model can be easily adapted to other settings through use of a different data loss. We manage the second assumption by use of virtual decimation (Groves et al., 2011; Simpson et al., 2012), which calculates the effective degrees of freedom of the residual map with stationary covariance by proxy of the correlation between neighbouring voxels in each direction.

We showed how good modelling choices, such as a careful choice of priors and transformation parametrisation can mitigate some of the issues caused by approximations necessary in practical applications. The main problem in inter-subject brain MRIregistration remains accuracy but the trade-off between the quality of the transformation and registration accuracy can be managed effectively.

The experiments were run on a system with an Intel i9-10920X CPU and a GeForce RTX 3090 GPU. Table 4 shows the runtime of the models. VItakes approximately $3\text{\,}\mathrm{min}$ to register a pair of images and produces 140 samples/sec, while SG-MCMCproduces 25 autocorrelated samples/sec. The registration time for SG-MCMCis absent because it is difficult to pin down the mixing time, so a fair comparison of the relative efficiency of VIand SG-MCMCcan be done only on the basis of the number of samples per second.

Due to lack of publicly available information we cannot directly compare the efficiency of our model to other Bayesian image registration methods. The speed of our model is an order of magnitude better than reported by Le Folgoc et al. (2017), while also being three- rather than two-dimensional. Thus, the proposed method based on SG-MCMCis very efficient given the Bayesian constraint. It is not as efficient as feed-forward neural networks, such as VoxelMorph. However, the proposed model is fully Bayesian and enables asymptotically exact sampling from the posterior of transformation parameters.