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

TL;DR

This paper introduces a memory-efficient Levenberg-Marquardt optimizer for diffeomorphic image registration, outperforming Adam in large-scale scenarios and demonstrating versatility across multiple medical imaging modalities.

Contribution

Proposes a novel factored Levenberg-Marquardt optimizer with minimal memory requirements, adaptable trust region tuning, and broad applicability to various image registration tasks.

Findings

Reduces memory usage by up to 24.6% for large volumes.

Matches or outperforms Adam on three out of four benchmarks.

Single hyperparameter configuration transfers across different modalities.

Abstract

FireANTs introduced a novel Eulerian descent method for plug-and-play behavior with arbitrary optimizers adapted for diffeomorphic image registration as a test-time optimization problem, with a GPU-accelerated implementation. FireANTs uses Adam as its default optimizer for fast and more robust optimization. However, Adam requires storing state variables (i.e. momentum and squared-momentum estimates), each of which can consume significant memory, prohibiting its use for significantly large images. In this work, we propose a modified Levenberg-Marquardt (LM) optimizer that requires only a single scalar damping parameter as optimizer state, that is adaptively tuned using a trust region approach. The resulting optimizer reduces memory by up to 24.6% for large volumes, and retaining performance across all four datasets. A single hyperparameter configuration tuned on brain MRI transfers without modification to lung CT and cross-modal abdominal registration, matching or outperforming Adam on three of four benchmarks. We also perform ablations on the effectiveness of using Metropolis-Hastings style rejection step to prevent updates that worsen the loss function.