ADAPTIVE GRADIENT DESCENT OPTIMIZATION OF INITIAL MOMENTA FOR GEODESIC SHOOTING IN DIFFEOMORPHISMS

Abstract

Diffeomorphic image registration algorithms are widely used in medical imaging, and require optimization of a high-dimensional nonlinear objective function. The function being optimized has many characteristics that are relevant for optimization but are typically not well understood. Due to that complexity, most authors have used a simple gradient descent, but it is not often discussed how step sizes are chosen or if line searches are used. Further, if a system is to be robust to a range of input images, that may differ to varying degrees, the optimization must be adaptable. Here, we present two methods of adaptable gradient descent with line searches, and test how they affect image registration. The optimization schemes are deployed for geodesic shooting in diffeomorphisms - an approach that is used to quantify anatomical changes, such as atrophy, in longitudinal image pairs. We evaluate the optimization schemes on their convergence characteristics and based on how well the resulting atrophy scores correlate with diagnostic group and mini mental state exam (MMSE) scores. We find that the Barzilai-Borwein method with a backtracking line search outperforms other optimization schemes in convergence time and adaptability by a wide margin. We also find that the variable optimization schemes do not significantly affect the ability to measure atrophy with clinical significance.

1. INTRODUCTION

The LDDMM (large deformation diffeomorphic metric mapping) framework for diffeomorphic image registration is described in significant procedural detail in the seminal paper by Beg et al. [1]. LDDMM proposes encoding the shape difference evident in two images of the same anatomy as a point on a manifold of diffeomorphisms. The objective of the algorithm is to construct a path on that manifold beginning at the identity and ending at the diffeomorphism that optimally matches the two images; Beg et al. show that at optimality the path is a geodesic. In the geodesic shooting formulation, the path is parameterized by an initial momentum vector field, from which the entire geodesic can be reconstructed by integrating the appropriate Euler-Poincaré differential equation (EPdiff) [2, 3, 4, 5, 6]. Also, Miller et al. [2] show that at optimality, the initial momentum vector field is proportional to the spatial gradient of the moving image. Hence, the objective of geodesic shooting in diffeomorphisms (GSiD) is to find a scalar momentum field that parameterizes an optimal matching between the given images.

An implementation of GSiD can be viewed as having three mathematical components: (1) construction of the model itself, which can include decisions about representation [6] and selection of parameters that dictate properties of the space of diffeomorphisms [7], (2) numerical integration of differential equations [3], and (3) an optimization procedure for the initial momentum field. The majority of work in the field has been in areas (1) and (2), with substantially less attention to paid to (3). With the exception of [5], most studies report using gradient descent with some step size ε, though details of the procedure including the determination of ε, are typically omitted. Very little information relevant to optimization is known about the GSiD objective function, such as its smoothness and curvature characteristics. Further, it is not known how variable these characteristics are to different inputs. Nonetheless, to build a GSiD system robust to variable inputs, some optimization procedure must be selected.

We consider here gradient descent with three different procedures to determine the step size ε: (1) a static step size, (2) a secant method line search, and (3) the Barzilai-Borwein method [8]. Methods (2) and (3) compute the step size at every iteration using limited local curvature data estimated from the objective function; hence, ε is adaptable to the inputs and the particular iteration of the optimization.

2. METHODS

2.1. GSiD

A complete discussion of the GSiD model is beyond the scope of this paper; for a thorough discussion of the following equations see [3]. For a moving image I and fixed image J, the GSiD objective function is:

$$E\left(P_0\right)=\frac{1}{{\sigma }^2}{\langle P_0\nabla I,K\left(P_0\nabla I\right)\rangle }_{L_2}+{\Vert I\circ {\varphi }_1^{-1}-J\Vert }^2$$ (1)

which must be minimized with respect to the initial scalar momentum field P₀. A given initial momentum provides the initial conditions for the EPdiff equation(s), which govern the time evolution of the momentum and moving image:

$$\{\begin{array}{l}{\partial }_{t}I+\nabla I\cdot v=0\hfill \\ {\partial }_{t}P+\nabla \cdot \left(P_v\right)=0\hfill \\ v+K\left(P\nabla I\right)=0\hfill \end{array}$$ (2)

The third equation states the relationship between momentum and velocity, where K plays the role of an inertia; K is a smoothing kernel and K (w) is taken to mean the convolution of vector field w with K. The path of diffeomorphisms ϕ_t is constructed from v(x, t) according to the ODE:

$$\{\begin{array}{l}{\partial }_t{\varphi }_t=v_t\left(\varphi \right)\hfill \\ {\varphi }_0=\text{Id}\hfill \end{array}$$ (3)

This yields the geodesic path of diffeomorphisms ϕ_t, where the end point ϕ₁ is used to match I and J. The matching residual $J-I\circ {\varphi }_1^{-1}$, which resides in the coordinate system of J, must be brought back to the coordinate system of I while respecting the geodesicity of the whole path ϕ_t. This is done by integrating the adjoint system backwards in time with initial conditions ${\widehat{I}}_1=J-I\circ {\varphi }_1^{-1}$ and ${\widehat{P}}_0=0$:

$$\{\begin{array}{l}{\partial }_t\widehat{I}+\nabla \cdot \left(v\widehat{I}\right)+\nabla \cdot \left(P\widehat{v}\right)=0\hfill \\ {\partial }_t\widehat{P}+v\cdot \nabla \widehat{P}+\nabla I\cdot \widehat{v}=0\hfill \\ \widehat{v}+K\left(\widehat{I}\nabla I-P\nabla \widehat{P}\right)=0\hfill \end{array}$$ (4)

${\widehat{P}}_0$ completes the gradient of equation (1) with respect to P₀. In a gradient descent scheme, that gradient is used to update P₀:

$$P_0^{k+1}=P_0^k-\epsilon \left(\nabla I\cdot K\left(P_0^k\nabla I\right)-{\widehat{P}}_0^k\right)$$ (5)

ε is one of the few user selected parameters in the GSiD model. A poor selection of ε can result in intractable compute times (if the user insists on running to convergence), sub-convergent results (if the optimization is stopped early due to time considerations), or numerical instability and divergence (if ε is too large). We are concerned in particular with the application of GSiD to longitudinal MRI time series of the brain to quantify atrophy. In that application, deformations are typically very low amplitude even relative to the spatial resolution of the images. Despite that, atrophy of as little as 5% in critical brain areas can have a significant impact on quality of life [9]. Hence, it is crucial to measure longitudinal deformations with the highest degrees of accuracy and precision possible. In such a case, the selection of ε can be critical to ensuring accurate and unbiased measurements.

2.2. Adaptable gradient descent steps

The simplest option is to select a priori a static value for ε which is fixed throughout the optimization. This value may perform well for some instances of data, and poorly for others. Even for a fixed input, it may perform well for a subset of iterations and poorly for others. We include this option as a baseline for comparison with more intelligent choices.

The secant line search method

Suppose we would like to optimize a function f(x) by gradient descent. Then, for iteration k, we would like to minimize f(x_k − ε_kf′(x_k)) with respect to ε_k. Take the truncated Taylor expansion (we temporarily omit the subscript k):

$$f\left(x-\epsilon f^{\prime }\right)\approx f\left(x\right)+\epsilon \left({\partial }_{\epsilon }f\left(x-\epsilon f^{\prime }\right){|}_{\epsilon =0}\right)+\frac{{\epsilon }^2}{2}\left({\partial }_{\epsilon }^{2}f\left(x-\epsilon f^{\prime }\right){|}_{\epsilon =0}\right)$$ (6)

If used directly, the second-order term will require the Hessian matrix of f. GSiD is a very high-dimensional optimization, hence the Hessian matrix f″ is intractable. We can replace the second-order term in the Taylor expansion with a finite difference approximation on the gradients:

$$\begin{gathered} {\partial }_{\epsilon }^{2}f\left(x-\epsilon f^{\prime }\right)\approx \frac{{\partial }_{\epsilon }f\left(x-\epsilon f^{\prime }\right){|}_{\epsilon =\sigma }-{\partial }_{\epsilon }f\left(x-\epsilon f^{\prime }\right){|}_{\epsilon =0}}{\sigma } \\ =\frac{-f^{\prime }{\left(x-\sigma f^{\prime }\right)}^T{f}^{\prime }+{f^{\prime }}^T{f}^{\prime }}{\sigma } \end{gathered}$$ (7)

Substitute (7) into (6), apply the remaining partial derivative, and further differentiate each side with respect to ε. You will arrive at the expression:

$${\partial }_{\epsilon }f\left(x-\epsilon f^{\prime }\right)\approx -{f^{\prime }}^T{f}^{\prime }+\frac{\epsilon }{\sigma }\left({f^{\prime }}^T{f}^{\prime }-f^{\prime }{\left(x-\sigma f^{\prime }\right)}^T{f}^{\prime }\right)$$

Finally, we set this equal to zero and solve for ε. Also, to use this formula for GSiD we must account for the metric in the space of momenta; the inner products must include the operator K. With these two final steps we arrive at the formula:

$${\epsilon }_k=\frac{{\sigma }_k{f^{\prime }}_k^{T}K\left({f^{\prime }}_k\right)}{{f^{\prime }}_k^{T}K\left({f^{\prime }}_k\right)-f^{\prime }{\left(x_k-{\sigma }_k{f}^{\prime }\right)}^{T}K\left({f^{\prime }}_k\right)}$$ (8)

Where ${f^{\prime }}_k=f^{\prime }\left(x_k\right)$. Essentially, the secant method approximates the objective function in the gradient direction as a parabola, the curvature of which is estimated by formula (7). Because the function may not be well estimated locally as a parabola, for a given gradient descent iteration k, formula (8) is applied iteratively giving a series of steps ${\epsilon }_k^i$. For the i^th secant method iteration, ${\sigma }_k^i=-{\epsilon }_k^{i-1}$, which leverages every gradient computation efficiently. ${\sigma }_k^0$ is set to a default value. This line search is stopped after a certain number of fixed iterations or when the magnitude of the update $\Vert \epsilon {f^{\prime }}_k\Vert$ falls below a threshold. Note, though we must evaluate multiple gradients during the line search iterations, we only move in the direction ${f^{\prime }}_k$ until the line search is stopped and we move to gradient descent iteration k + 1.

The Barzilai-Borwein method

We derive the method assuming $f\left(x\right)=\frac{1}{2}{x}^{T}Ax-b^{T}x$. The gradient is then f′ = Ax − b and the Hessian is f″ = A. Newton’s method, a second-order optimization that accounts for the curvature of the objective function, proceeds as $x_{k+1}=x_k-A^{-1}{f^{\prime }}_k$. (For a symmetric positive definite quadratic form, this will converge in a single step and is equivalent to Gaussian elimination). The objective of the Barzilai-Borwein (BB) method is to let ε be determined by the simplest possible approximation to Newton’s method:

$$-\epsilon f^{\prime }=-{\left({\epsilon }^{-1}\text{Id}\right)}^{-1}{f}^{\prime }\approx -A^{-1}{f}^{\prime }$$ (9)

Let s_k = x_k − x_k−1 and $y_k={f^{\prime }}_k-{f^{\prime }}_{k-1}$. For the quadratic form, A satisfies As_k = y_k. So, we will let ε be the solution to the least squares problem:

$${\epsilon }_k={\text{argmin}}_{\alpha }\frac{1}{2}{\Vert s_k-\alpha y_k\Vert }^2$$ (10)

which has the closed-form solution:

$${\epsilon }_k=\frac{s_k^T{y}_k}{y_k^T{y}_k}$$ (11)

Again, to apply this to GSiD we must account for the metric in the space of momenta:

$${\epsilon }_k=\frac{s_k^{T}K\left(y_k\right)}{y_k^{T}K\left(y_k\right)}$$ (12)

Similar to the secant method, the BB method approximates second-order information, but it does not require a second gradient computation. Even so, in some places formula (9) is likely to be a very poor approximation. It is well known that as a result, BB step sizes do not provide monotonic optimization; that is, occasionally ε_k is too large. However, for nonlinear optimization, some degree of nonmonotonicity may be desired as it may help escape spurious local minima. Hence, the BB method is often coupled with a backtracking line search [10]. In our case, ε_k is iteratively cut in half until the first Wolfe condition is satisfied:

$$f\left(x_k-{\epsilon }_k{f^{\prime }}_k\right)\le {max}_{j}f\left(x_j\right)-\gamma {\epsilon }_k{f^{\prime }}_k^{T}K\left({f^{\prime }}_k\right)$$ (13)

where max_(k−M,₀₎ ≤ j ≤ k. M controls the degree of monotonicity (we use M =10) and γ is related to our expectation of the objective function’s local curvature. γ is typically chosen to be small (we use γ = 10⁻⁴).

3. EXPERIMENTAL RESULTS

We took 100 randomly chosen subjects from the ADNI-2 longitudinal MRI dataset - available at adni.loni.usc.edu - and registered their baseline scans to their 24-month follow-up scans using GSiD. We did these registrations under 5 different experimental conditions: static step sizes of 0.001, 0.01, and 0.1, the secant method with ${\epsilon }_k^0=0.01$, and the BB method with ε₀ = 0.01. For all approaches, the optimization was stopped when the gradient magnitude (relative to the initial gradient magnitude) fell below a chosen threshold, or after 300 iterations, whichever came first. Before GSiD, the images were preprocessed according to the protocol detailed in [11, 12].

After GSiD, the Jacobian determinants of the deformations mapping the baseline to the 24 month followup images were moved to a common coordinate system. The Jacobian determinants were averaged in a region where the rate of atrophy is significantly associated with Alzheimer’s Disease (AD), a stat-ROI, (Fig. 1) to produce a scalar value atrophy score that represents the percent volume loss within the region for each subject [9]. The region was constructed from a non-overlapping data set from that evaluated here.

Fig. 1.

Region of interest with significant atrophy in AD, used here to compute atrophy scores

Figures 2 and 3 show convergence characteristics of the five optimizations strategies. Contrary to equation (1), we did not use sum of squared differences to drive the registration. We used the squared Local Correlation Coefficient (LCC) which is also used in [13]; LCC increases as the images become better matched. Curves that do not extend the full 300 iterations are instances that stopped early due to the gradient magnitude stopping criterion. The largest static step size clearly causes oscillations in all instances. The smallest static step size did not permit any instances to complete before reaching 300 iterations, it is likely that many instances are sub-convergent. The middle static step size appears to be a good compromise, but for many instances, the gradient magnitude oscillates. Apparently, none of the static step sizes is appropriate for all instances of the data or through all iterations of the optimization. The secant and BB methods show better convergence characteristics, with more instances finishing early. However, for the secant method not all instances converged. The spikes in the gradient magnitude for the BB method are due to the nonmonotonicity discussed above.

Fig. 2.

Convergence of data match term; curves that do not extend the full 300 iterations stopped early due to the gradient magnitude stopping criteria.

Fig. 3.

Convergence of gradient magnitude; curves that do not extend the full 300 iterations stopped early due to the gradient magnitude stopping criteria.

A good first question to ask is whether the choice of optimization procedure had a significant impact on atrophy scores. Figure 4 shows the p-values from paired t-tests between the measured atrophy scores for all pairs of optimization approaches. All five optimization procedures produced atrophy measurements that were significantly different from the others. Figure 3 also shows the average number of iterations and the number of instances that failed to converge due to numerical instability. In practice, the failed instances would have to be rerun with the parameters adjusted by hand. The BB method clearly had the fastest convergence, and was sufficiently adaptable that no instances failed to converge.

Fig. 4.

Statistical tests, convergence information, and correlations; DX: diagnostic group; MMSE: Mini Mental State Exam

Atrophy measurements such as these have been shown to correlate with diagnostic category and performance on cognitive tests. The data set included subjects from four diagnostic categories: healthy controls (HC), early mild cognitive impairment (eMCI), late mild cognitive impairment (lMCI), and Alzheimer’s disease (AD). Each subject also had a mini mental state exam (MMSE) administered at the 24 month follow up time point. The MMSE scores from 0 - 30, where scores below 24 typically indicate some level of dementia. Figure 4 also shows Pearson’s correlation coefficients between atrophy scores and diagnostic group and also between atrophy scores and MMSE scores for each of the five optimization approaches. The correlations appear sufficiently similar across optimization approaches to suggest that faster or more adaptable optimization approaches do not compromise the ability to measure clinically meaningful atrophy.