Factored Levenberg-Marquardt for Diffeomorphic Image Registration: An efficient optimizer for FireANTs
Authors: Rohit Jena, Pratik Chaudhari, James C. Gee
Organizations: Department of Computer and Information Science, University of Pennsylvania, Philadelphia, PA 19104, USA · Department of Electrical and Systems Engineering, University of Pennsylvania, Philadelphia, PA 19104, USA · Department of Radiology & Penn Image Computing and Science Laboratory (PICSL), Perelman School of Medicine, University of Pennsylvania, Philadelphia, PA 19104, USA
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.
Figures & tables
Dataset
Method
Mean
Std
LUMIR (Dice ↑ )
Adam
0.8650
0.0201
Levenberg-Marquardt
0.8715
0.0163
OASIS (Dice ↑ )
Adam
0.8028
0.0333
Levenberg-Marquardt
0.8293
0.0313
NLST (TRE ↓ )
Adam
0.85
0.28
Levenberg-Marquardt
1.68
2.16
Table 1: Optimizer comparison across four registration benchmarks. Mean Dice ( ↑ ) for LUMIR, OASIS, Abdomen-L2R; mean TRE in mm ( ↓ ) for NLST. “Pairs won” counts image pairs where the method beats Adam. Best result per dataset in bold .
Label
Adam
LM
Spleen
0.8344
0.8325
Right kidney
0.7288
0.7211
Left kidney
0.7594
0.7662
Liver
0.7208
0.7262
Table 2: Per-label Dice on Abdomen-L2R (4 organs, 48 pairs). LM performs similarly to Adam across all labels.
Figure 1: Sensitivity of mean Dice on LUMIR (38 pairs) to all three LM hyperparameters. Our defaults ( λ0=0.006 , μ+=1.5 , μ−=0.975 , red dashed) were selected via Bayesian optimization on LUMIR and are marked in each panel. Left : λ0 is entirely insensitive—Dice varies by <0.001 over a 5000× range. Center : μ+ has a catastrophic cliff between 2.3 and 3.0 ; our default of 1.5 sits at the safe peak. Right : μ− shows monotone improvement up to 0.963 then a small drop; our default of 0.975 sits in the optimal plateau.
Figure 2: Mean Dice ( ± std) on LUMIR (38 pairs) as tile size k varies from 1 to 10, with all other LM hyperparameters fixed to their defaults. Pointwise updates ( k=1 , red dashed) achieve the best mean Dice ( 0.8716 ) and lowest variance; performance drops −0.009 at k=2 and degrades monotonically to 0.828 by k=7 – 10 , with standard deviation rising throughout.
Dataset
Adam
LM-transfer
LM-tuned
LUMIR (Dice)
0.8650
—
0.8715
OASIS (Dice)
0.8028
0.8289
0.8293
NLST (TRE, mm)
0.85
1.68 ‡
1.68
Abdomen (Dice)
0.7609
0.7615
0.7615
‡ Inflated by one outlier pair; excluding it: 0.95 mm vs. 0.83 mm for Adam.
Table 3: Cross-dataset transfer of LUMIR-tuned LM hyperparameters compared to per-dataset Optuna search. “LM-transfer” applies LUMIR parameters directly with no additional tuning; “LM-tuned” uses dataset-specific search.
Dataset
Condition
Mean
Std
LUMIR (Dice ↑ )
No rejection
0.8715
0.0163
+ rejection, cap =∞
0.8654
0.0303
+ rejection, cap =1.0
0.8709
0.0170
OASIS (Dice ↑ )
No rejection
0.8028
0.0333
+ rejection, cap =∞
0.7362
0.0899
+ rejection, cap =1.0
0.8014
0.0343
Table 4: Effect of step rejection on default LM. All conditions use τ=1.0 . Without a cap, rejection is consistently harmful across all datasets. With λmax=1.0 , rejection is dataset-dependent: neutral on LUMIR and OASIS, mildly harmful on Abdomen-L2R (48 pairs), and beneficial on NLST by suppressing one catastrophic outlier. Best result per dataset in bold .
Peak memory (MB)
Time/step (s)
N
Adam
LM
Mem. saved
Adam
LM
Speedup
64
61.1
55.9
8.6%
0.593
0.356
1.67×
128
433.1
384.9
11.1%
0.605
0.606
1.00×
192
1444.1
1282.7
11.2%
1.085
1.062
1.02×
256
3400.1
3022.1
11.1%
2.127
2.061
1.03×
320
6638.1
5898.0
11.1%
3.937
3.801
1.04×
Table 5: Peak GPU memory (MB) and time per step (s) on an NVIDIA A6000, for cubic volumes of side N . LM is never slower and consistently uses less memory. Best result in bold .
Medical image registration aligns scans across time, patients, or imaging modalities to track anatomy and support diagnosis, treatment, and image-guided procedures.
Department of Brain Sciences, Imperial College London, London, UK · Department of Computing, Imperial College London, London, UK · School of Computer Science, University of Birmingham, Birmingham, UK +2
Department of Information Technology, Uppsala University, Uppsala, Sweden · Department of Medical Physics, Uppsala University Hospital, Uppsala, Sweden · Department of Immunology, Genetics and Pathology, Uppsala University, Uppsala, Sweden