Source-linked AI summary
HyperMorph: Amortized Hyperparameter Learning for Image Registration
Andrew Hoopes, Malte Hoffmann, Bruce Fischl, John Guttag, Adrian V. Dalca
TL;DR
Registration performance depends on hyperparameters, but conventional tuning repeatedly trains models across candidate values and can be computationally and labor intensive. HyperMorph learns a hypernetwork that generates registration-network parameters for continuous hyperparameter inputs, enabling rapid test-time tuning. The approach reduces compute, improves robustness to initialization, and supports optimization for datasets, tasks, and anatomical regions without retraining.
Problem
Registration quality depends strongly on hyperparameters, while conventional tuning evaluates separate models across values and can require extensive computation and human effort.
Method
HyperMorph learns a hypernetwork that maps loss hyperparameters to registration-network parameters, representing a continuous landscape of registration models with one rich model.
Results
HyperMorph requires 5.2 ±0.2 times fewer GPU-hours for single-hyperparameter tests and 10.5 ± 0.2 times fewer for two-hyperparameter models than baseline grid search.
Takeaways & Limitations
A single HyperMorph model enables rapid, fine-grained hyperparameter discovery for different datasets, registration tasks, and anatomical regions without retraining.
Abstract
from arXiv · showhide
We present HyperMorph, a learning-based strategy for deformable image registration that removes the need to tune important registration hyperparameters during training. Classical registration methods solve an optimization problem to find a set of spatial correspondences between two images, while learning-based methods leverage a training dataset to learn a function that generates these correspondences. The quality of the results for both types of techniques depends greatly on the choice of hyperparameters. Unfortunately, hyperparameter tuning is time-consuming and typically involves training many separate models with various hyperparameter values, potentially leading to suboptimal results. To address this inefficiency, we introduce amortized hyperparameter learning for image registration, a novel strategy to learn the effects of hyperparameters on deformation fields. The proposed framework learns a hypernetwork that takes in an input hyperparameter and modulates a registration network to produce the optimal deformation field for that hyperparameter value. In effect, this strategy trains a single, rich model that enables rapid, fine-grained discovery of hyperparameter values from a continuous interval at test-time. We demonstrate that this approach can be used to optimize multiple hyperparameters considerably faster than existing search strategies, leading to a reduced computational and human burden as well as increased flexibility. We also show several important benefits, including increased robustness to initialization and the ability to rapidly identify optimal hyperparameter values specific to a registration task, dataset, or even a single anatomical region, all without retraining the HyperMorph model. Our code is publicly available at http://voxelmorph.mit.edu.
1 Introduction
HyperMorph addresses costly registration hyperparameter tuning by learning a single model that represents registration networks across continuous hyperparameter values. This enables faster tuning, greater robustness to initialization, and task- or region-specific optimization without retraining.
- Registration quality depends strongly on hyperparameters, whose optimal values can vary across modalities, anatomy, datasets, and tasks.
- Traditional grid or random search repeatedly trains separate registration models for discrete hyperparameter values, requiring substantial computational and human effort.
- HyperMorph learns a continuous hyperparameter-to-registration mapping with one rich model, approximating a landscape of models instead of retraining each configuration.
- HyperMorph is more robust to initialization than standard registration models, reducing the need to retrain models with different initializations.
- A single HyperMorph model supports rapid, precise hyperparameter discovery for different datasets, anatomical regions, and registration tasks at test time.
2 Related Work
Image registration methods use hyperparameters to balance image matching and regularization, but tuning them is expensive. Existing optimization strategies reduce search inefficiency in different ways, while hypernetworks provide a route to differentiable, conditioned model weights.
- Image Registration: Classical registration methods independently optimize deformation fields for each image pair, whereas learning-based methods use CNNs to rapidly predict deformation fields.
- Image Registration: Registration methods commonly use hyperparameters to balance image matching against regularization or smoothness, with additional parameters appearing in loss terms.
- Image Registration: Classical pairwise registration can take tens of minutes or more, while learning-based model training can require days and make hyperparameter search consume hundreds of GPU-hours.
- Hyperparameter Optimization: Grid, random, sequential, Bayesian, bandit-based, and gradient-based optimization methods seek efficient hyperparameter search but still involve repeated or costly model evaluations.
- Hypernetworks: Hypernetworks output the weights of a primary network and enable differentiation through the model with respect to selected hyperparameters.
3 Methods
HyperMorph learns a hypernetwork that maps registration hyperparameters to registration-network parameters, enabling one model to produce deformation fields across a continuous hyperparameter range. The framework supports unsupervised and semi-supervised losses, interactive or automatic tuning, and task-specific registration settings without retraining.
- HyperMorph framework: HyperMorph defines a registration network that maps moving and fixed images to a dense deformation field φ.The network is written as gθg(m, f) = φ.
- HyperMorph framework: The hypernetwork hθh(Λ) outputs registration-network parameters θg from sampled loss hyperparameters Λ.HyperMorph learns the hypernetwork parameters using stochastic gradient methods.
- HyperMorph framework: HyperMorph trains a single model to approximate registration networks continuously across a range of hyperparameter values.The hypernetwork is trained end-to-end and exploits implicit weight-sharing across that landscape.
- Training objective: Training samples hyperparameters from a prior distribution and uses them both as hypernetwork inputs and in the registration loss.The prior can be uniform over a predefined range or adapted from prior expectations.
- Training objective: The unsupervised loss combines image similarity and deformation regularity, with λ, λsim, and λreg controlling their relative contributions.Similarity metrics include MSE, local NCC, or NMI depending on image modality; local NCC and NMI have additional hyperparameters.
- Hyperparameter tuning: A trained HyperMorph model supports interactive slider-based tuning or automatic validation-based optimization by treating Λ as learnable while freezing the hypernetwork.Interactive tuning can select settings for specific anatomical regions, while automatic tuning uses validation segmentation maps.
4 Experiments
The experiments evaluate HyperMorph’s efficiency, capacity, robustness, and tuning utility against baseline grid-search models. A single model captures hyperparameter-dependent registration behavior while reducing computation and supporting task-specific optimization without retraining.
- Efficiency and capacity: A single HyperMorph model captures the behavior of registration models trained with separate hyperparameter values.The evaluation tests whether one model can represent a landscape of baseline models across registration settings.
- Efficiency and capacity: HyperMorph requires 5.2 ± 0.2 times fewer GPU-hours for one hyperparameter and 10.5 ± 0.2 times fewer for two hyperparameters than baseline grid searches.The two-hyperparameter experiments use 2D models because the corresponding 3D grid search is computationally prohibitive.
- Efficiency and capacity: The optimal λ difference is 0.04 ± 0.02 with a maximum Dice difference of 0.16 ± 0.03 in single-hyperparameter experiments, while multi-hyperparameter experiments show a maximum Dice difference of 0.32 ± 0.02.These results indicate similar performance to dense baseline grids despite the computational advantage.
- Robustness to initialization: HyperMorph has lower initialization variability than baselines, with average Dice standard deviations 2.79 times lower for MSE and 2.16 times lower for NMI.The authors associate this robustness with a lower likelihood of converging to local minima and reduced need for repeated training.
- Hyperparameter-tuning utility: Automatic optimization takes 2.8 ± 0.3 minutes using 20 validation pairs and identifies λ values that vary across subpopulations, tasks, and anatomical regions.The experiments also demonstrate interactive tuning and optimization without retraining the HyperMorph model.
5 Conclusion
HyperMorph uses a hypernetwork to predict registration-network parameters for chosen hyperparameters, enabling efficient tuning without repeatedly training separate models. It supports rapid optimization across datasets, tasks, and anatomical regions, while revealing region-specific regularization preferences.
- HyperMorph predicts registration-network parameters from desired hyperparameters, replacing repeated model training with one tunable model.
- Optimal regularization weights λ∗ vary across anatomical labels, indicating that downstream registration tasks may require region-specific choices.
- A single HyperMorph model enables fast discovery of optimal hyperparameters for different datasets, registration tasks, and individual anatomical regions.