Factored Levenberg-Marquardt for Diffeomorphic Image Registration: An efficient optimizer for FireANTs

1. Problem Statement

The paper addresses the computational and memory bottlenecks associated with diffeomorphic image registration when using modern deep learning frameworks like FireANTs.

• Context: FireANTs frames registration as a test-time optimization problem on GPUs, minimizing a loss function $L(u)$ consisting of image similarity and regularization terms.
• Current Limitation: The default optimizer, Adam, is a quasi-second-order method that maintains exponential moving averages of gradients (momentum) and squared gradients. While effective, these state variables consume significant GPU memory (proportional to the number of voxels), making Adam impractical for very large 3D medical volumes (e.g., high-resolution CT or MRI).
• Goal: Develop a more memory-efficient second-order optimizer that retains the robustness and speed of Adam without the heavy memory overhead of storing moment buffers.

2. Methodology: Factored Levenberg–Marquardt (LM)

The authors propose a modified Levenberg–Marquardt (LM) optimizer tailored for dense displacement fields.

A. Core Innovation: Per-Voxel Rank-1 Factorization

Standard LM requires solving a system involving the Hessian matrix $J^\top J$. For a 3D displacement field with $\sim 10^7$ parameters, the full Hessian is intractable ($\sim 10^{14}$ parameters).

• Approximation: The authors observe that the registration gradient at a specific voxel depends primarily on local image content. They approximate the Gauss-Newton Hessian independently per voxel using a rank-1 outer product: $\hat{H}_x = g_x g_x^\top$.
• Closed-Form Update: This allows for a closed-form solution for the update step $\Delta u_x$ without matrix inversion or storage:
  $$\Delta u_x = -\frac{r g_x}{\|g_x\|^2 + \lambda}$$
  where $r$ is the scalar residual, $g_x$ is the gradient at voxel $x$, and $\lambda$ is the damping parameter.
• Tile Size: While the method can pool gradients over spatial tiles (larger windows) to form a $d \times d$ matrix, experiments show that the pointwise update (tile size = 1) yields the best performance and lowest cost.

B. Adaptive Damping Strategy

The damping parameter $\lambda$ is crucial for balancing between Newton's method (fast near optima) and gradient descent (robust far from optima).

• Trust Region Heuristic: $\lambda$ is updated adaptively after every step based on the change in loss:
  • If the loss increases (bad step): $\lambda \leftarrow \mu_+ \cdot \lambda$ (increase damping).
  • If the loss decreases (good step): $\lambda \leftarrow \mu_- \cdot \lambda$ (decrease damping).
• Rejection Criterion: To prevent steps that significantly worsen the loss, a Metropolis-Hastings-style rejection step is implemented. If the loss degradation exceeds a threshold relative to the previous improvement, the step is rejected, $\lambda$ is increased, and the step is retried.

C. Integration with FireANTs

The LM optimizer is integrated into FireANTs' Eulerian descent framework. It updates the velocity field $v(t)$, which is then composed with the current warp $u(t)$ to maintain diffeomorphism, bypassing the need for expensive parallel transport operations.

3. Key Contributions

1. Memory Efficiency: The proposed LM optimizer requires only a single scalar state variable ($\lambda$) per iteration, eliminating the need for momentum and squared-momentum buffers.
2. Second-Order Behavior at First-Order Cost: By utilizing rank-1 Hessian factorization, the method achieves second-order convergence properties with computational costs comparable to Adam.
3. Hyperparameter Universality: The authors demonstrate that a single hyperparameter configuration (tuned on one dataset) transfers effectively across different imaging modalities (MRI, CT) and anatomical regions (Brain, Lung, Abdomen) without re-tuning.
4. Sensitivity Analysis: The paper identifies a critical "damping cliff" where the increase factor $\mu_+$ must be kept below 2.3; exceeding this causes the optimizer to over-damp and stall.

4. Experimental Results

Experiments were conducted on four benchmarks: LUMIR and OASIS (Brain MRI), NLST (Lung CT), and Abdomen-L2R (Cross-modal MR-to-CT).

• Performance Accuracy:
  • Brain MRI (LUMIR/OASIS): LM outperformed Adam significantly (e.g., +0.027 Dice on OASIS).
  • Abdomen (MR-CT): LM matched Adam's performance.
  • Lung CT (NLST): Adam initially appeared better, but this was driven by a single outlier. When the outlier was excluded, LM actually outperformed Adam (0.95 mm vs. 0.83 mm TRE).
• Memory Savings:
  • LM reduced peak GPU memory usage by 11% to 24.6% depending on volume size.
  • For a $512^3$ volume, LM saved approximately 3,024 MB of VRAM compared to Adam.
• Runtime:
  • LM is never slower than Adam.
  • For small volumes, LM is up to 1.67× faster because it avoids the overhead of updating moment buffers.
  • For large volumes, it is roughly 4% faster (compute-bound).
• Hyperparameter Transfer: A configuration tuned on LUMIR ($\lambda_0=0.006, \mu_+=1.5, \mu_-=0.975$) transferred directly to all other datasets with no modification, confirming the optimizer's universality.

5. Significance and Implications

• Enabling Large-Scale Registration: By reducing memory consumption by up to 25%, this method enables diffeomorphic registration on significantly larger 3D volumes that were previously infeasible with Adam due to GPU memory constraints.
• Rehabilitating Second-Order Methods: The paper demonstrates that second-order methods can be made practical for deep learning-based registration if the Hessian is approximated efficiently (rank-1) and damping is managed carefully.
• Robustness: The "damping cliff" analysis provides critical guidance for practitioners, warning that aggressive damping increases ($\mu_+ > 2.3$) can catastrophically stall optimization.
• Practical Adoption: The authors provide a "universal" set of hyperparameters that work across modalities, removing the need for dataset-specific tuning and simplifying the deployment of high-performance image registration pipelines.

In conclusion, the Factored Levenberg–Marquardt optimizer offers a highly efficient, memory-light, and robust alternative to Adam for diffeomorphic image registration, making it particularly suitable for large-scale medical imaging applications.