Factored Levenberg-Marquardt for Diffeomorphic Image Registration: An efficient optimizer for FireANTs
This work addresses memory limitations for large-scale image registration tasks, offering an efficient alternative to Adam with minimal hyperparameter tuning, though it is incremental as it builds on existing FireANTs methods.
The paper tackles the high memory consumption of Adam optimizer in diffeomorphic image registration by proposing a modified Levenberg-Marquardt optimizer with a single scalar damping parameter, reducing memory by up to 24.6% for large volumes while matching or outperforming Adam on three of four benchmarks.
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.