Fast Nonparametric Mutual-Information-based Registration and Uncertainty Estimation

Abstract

In this paper we propose a probabilistic model for multi-modal non-linear registration that directly incorporates the mutual information (MI) metric into a demons-like optimization scheme. In contrast to uni-modal registration, where the demons algorithm uses repeated spatial filtering to obtain very fast solutions, MI-based registration currently relies on general-purpose optimization schemes that are much slower. The central idea of this work is to reformulate an often-used histogram interpolation technique in MI implementations as an explicit spatial interpolation step within a generative model. By exploiting the specific structure of this model, we obtain a dedicated and fast expectation-maximization optimizer with demons-like properties. This also leads to an easy-to-implement Gibbs sampler to infer registration uncertainty in high-dimensional deformation models, involving very little additional code and no external tuning. Preliminary experiments on multi-modal brain MRI images show that the proposed optimizer can be both faster and more accurate than the free-form deformation method implemented in Elastix. We also demonstrate the sampler’s ability to produce direct uncertainty estimates of MI-based registrations – to the best of our knowledge the first method in the literature to do so.

Appendix: Connection with MI-based registration

MI-based registration with partial volume interpolation can be interpreted as implicitly using the proposed generative model but with a different optimization strategy, in which EM is used to estimate $\varvec{\theta }$ but not $\mathbf {d}$. In particular, $\mathbf {\widehat{d}}$ can also be estimated by optimizing $\log p( \mathbf {d}, \varvec{\widehat{\theta }}_d| \mathbf {u})$ for $\mathbf {d}$ with a general-purpose optimizer, where $\varvec{\widehat{\theta }}_d= \arg \max _{ \varvec{\theta }} \log p( \mathbf {d}, \varvec{\theta }| \mathbf {u})$ involves an inner optimization that for each $\mathbf {d}$ estimates a matched $\varvec{\widehat{\theta }}_d$ de novo from starting values $\theta _{k,l}^0 = 1/L, \forall k,l$ by interleaving the EM Eqs. (2) and (3). When a flat prior $p(\varvec{\theta }) \propto 1$ is used, the resulting effective registration criterion is then directly related to MI as follows:

$$\begin{aligned} \log p( \mathbf {d}, \varvec{\widehat{\theta }}_d| \mathbf {u}) \simeq I \mathrm {MI}( \mathbf {d}) + \log p( \mathbf {d}) + \mathrm {const}, \end{aligned}$$

(7)

where

$$\begin{aligned} \mathrm {MI}( \mathbf {d}) = \sum _{k=1}^K \sum _{l=1}^L n_{k,l} \log \frac{ n_{k,l} }{ n_k \, n_l } \quad \mathrm {with} \quad n_{k,l} = \frac{N_{k,l}}{I}, \,\, n_k = \sum _{l} n_{k,l}, \,\, n_l = \sum _k n_{k,l} \end{aligned}$$

is the MI criterion using partial volume interpolation [18, 19], in which joint histogram counts $N_{k,l}$ are computed from fractional weights $\bar{w}_{i,j}^d = \beta ^b\left( y_j - (x_i + d_i ) \right)$ as in Eq. (4). To see why Eq. (7) holds, we can also write $\mathrm {MI}( \mathbf {d})$ as

$$\begin{aligned} \mathrm {MI}(\mathbf {d}) = \frac{1}{I} \sum _{i=1}^I \sum _{j=1}^J \bar{w}_{i,j}^d \log \bar{\theta }_{v_j,u_i}^d - \underbrace{\sum _{l=1}^L n_l \log n_l}_{\mathrm {const}} \quad \mathrm {with} \quad \bar{\theta }_{kl}^d = n_{k,l}/n_k , \end{aligned}$$

(8)

and, since $\log p( \mathbf {d}, \varvec{\widehat{\theta }}_d) = Q( \mathbf {d}, \varvec{\widehat{\theta }}_d| \mathbf {d}, \varvec{\widehat{\theta }}_d)$,

$$\begin{aligned} \log p( \mathbf {d}, \varvec{\widehat{\theta }}_d| \mathbf {u}) - \log p(\mathbf {d})= & {} \sum _{i=1}^I \sum _{j=1}^J \widehat{w}_{i,j}^d \log \widehat{\theta }_{v_j,u_i}^d \nonumber \\&- \sum _i D_{KL}\left[ p( n_i | u_i, \mathbf {d}, \varvec{\widehat{\theta }}_d) \, \Vert \, p( n_i | \mathbf {d}) \right] + \mathrm {const}, \end{aligned}$$

(9)

where $\widehat{w}_{i,j}^d = w_{i,j}( \mathbf {d}, \varvec{\widehat{\theta }}_d)$ and $D_{KL}( . \Vert . )$ denotes the Kullback-Leibler (KL) divergence. Comparing Eqs. (8) and (9), and noting that $\bar{w}_{i,j}^d$ and $\varvec{\bar{\theta }}_d$ are precisely the weights and estimate of $\varvec{\theta }$ in the first iteration of the inner EM optimization, MI-based registration can therefore be interpreted as making a “lazy” attempt at measuring $\log p( \mathbf {d}, \varvec{\widehat{\theta }}_d)$, using only a single iteration in the inner optimization of $\varvec{\widehat{\theta }}_d$, and ignoring the KL divergence between the prior and the posterior node assignment distributions. In the special case where $p( n_i | \mathbf {d})$ takes only binary values $\{0,1\}$, the approximation in Eq. (7) will be exact since the EM algorithm then immediately finds $\varvec{\widehat{\theta }}_d$ in its first iteration and the KL divergence term vanishes. This will happen when B-splines of order $b=0$ are used, or for first-order B-splines ($b=1$) whenever the image grids of $\mathbf {u}$ and $\mathbf {v}$ perfectly align.