Source-linked AI summary
An Unsupervised Learning Model for Deformable Medical Image Registration
Guha Balakrishnan, Amy Zhao, Mert R. Sabuncu, John Guttag, Adrian V. Dalca
TL;DR
Existing registration methods optimize each image pair separately, making deformable registration slow for large datasets. This paper learns a shared CNN-based registration function without supervised labels and evaluates it directly on new scan pairs. The method achieves comparable accuracy to state-of-the-art 3D registration while reducing runtime by orders of magnitude.
Problem
Per-pair deformable registration optimization is computationally intensive and extremely slow in practice.
Method
VoxelMorph learns shared CNN parameters from volume pairs and obtains a new registration field by evaluating the learned function without ground-truth fields or landmarks.
Results
The method achieves comparable registration accuracy to state-of-the-art 3D registration while operating orders of magnitude faster.
Takeaways & Limitations
The approach can significantly speed up medical image analysis and processing pipelines while supporting accuracy–runtime tradeoffs through model design.
Takeaways & Limitations
Unlike ANTs, the models are not guaranteed to produce diffeomorphic registrations with invertibility and topology preservation.
Abstract
from arXiv · showhide
We present a fast learning-based algorithm for deformable, pairwise 3D medical image registration. Current registration methods optimize an objective function independently for each pair of images, which can be time-consuming for large data. We define registration as a parametric function, and optimize its parameters given a set of images from a collection of interest. Given a new pair of scans, we can quickly compute a registration field by directly evaluating the function using the learned parameters. We model this function using a convolutional neural network (CNN), and use a spatial transform layer to reconstruct one image from another while imposing smoothness constraints on the registration field. The proposed method does not require supervised information such as ground truth registration fields or anatomical landmarks. We demonstrate registration accuracy comparable to state-of-the-art 3D image registration, while operating orders of magnitude faster in practice. Our method promises to significantly speed up medical image analysis and processing pipelines, while facilitating novel directions in learning-based registration and its applications. Our code is available at https://github.com/balakg/voxelmorph .
1. Introduction
Deformable registration is important but traditionally slow because each image pair requires computationally intensive optimization. The proposed VoxelMorph approach learns a shared CNN-based registration function without supervised labels, enabling rapid registration with comparable accuracy.
- Traditional deformable registration establishes dense nonlinear correspondences while enforcing smoothness constraints.
- VoxelMorph replaces per-pair optimization with one global optimization of a parametrized registration function learned from a volume collection.
- The CNN uses shared parameters across a population so new volume pairs can be registered by evaluating the learned function.
- The method requires no supervised information, including ground-truth correspondences or anatomical landmarks.
- The method is presented as broadly applicable to registration tasks within and beyond medical imaging.
- Scans that previously required two hours can be registered in one or two minutes on a CPU and under a second on a GPU, with comparable accuracy to state-of-the-art registration.
2. Background
Medical deformable registration warps one volume to another using a dense nonlinear field after affine alignment. Standard formulations balance image similarity against deformation smoothness, while anatomical variability motivates the deformable step.
- Deformable registration computes dense nonlinear voxel correspondences after an initial affine transformation for global alignment.
- Existing algorithms typically optimize a transformation independently for each volume pair using an energy function.
- The registration objective combines image-similarity loss with a smoothness penalty on the registration field, weighted by λ.
- Common similarity metrics include mean squared voxel difference, mutual information, and cross-correlation, with the latter two useful for varying intensity distributions and contrasts.
- MRI scans show significant structural differences across subjects, motivating deformable registration for analyzing interscan variation.
3. Related Work
Prior registration work includes displacement-field, diffeomorphic, optical-flow, and neural approaches. Learning-based methods can avoid iterative optimization, but many require supervised registrations or remain limited in demonstrated scope.
- Traditional 3D registration includes elastic models, statistical parametric mapping, free-form B-splines, and Demons methods based on displacement fields.
- Diffeomorphic transforms preserve topology and enforce invertibility, distinguishing them from displacement-field approaches.
- Earlier unsupervised neural registration methods used CNNs and spatial transformation functions but were demonstrated only on limited settings and had significant drawbacks.
- Optical flow is an analogous 2D problem that estimates dense displacement fields using variational, feature-based, or nearest-neighbor methods.
- Several neural optical-flow and image-alignment methods learn parametric CNN functions but require ground-truth registrations during training.
4. Method
VoxelMorph learns a CNN-based registration function from aligned 3D volume pairs and applies it to unseen pairs without supervised registration labels. A differentiable spatial transformation reconstructs the moving image, while similarity and smoothness losses train anatomically plausible deformation fields.
- Registration formulation: The method assumes affine preprocessing has removed global misalignment, leaving nonlinear registration between grayscale volumes.The paper focuses on n = 3 and affinely aligned inputs.
- Registration function: A CNN gθ(F, M) outputs a registration field φ mapping voxels in the fixed volume to similar anatomical locations in the moving volume.The network parameters θ are learned convolutional kernels.
- Learning procedure: Training minimizes an expected loss over volume pairs, and testing obtains a field by directly evaluating g on an unseen pair without ground truth fields or landmarks.The shared parameters are optimized with stochastic gradient descent.
- CNN architecture: The encoder-decoder CNN uses skip connections and hierarchical 3D features to estimate dense correspondences, with variants trading registration accuracy against computation time.VoxelMorph-2 adds an output-resolution convolution and more late-layer channels than VoxelMorph-1.
- Spatial transformation: The spatial transformer warps M using φ with differentiable interpolation, allowing image-similarity errors to backpropagate through the registration operation.Each subpixel location is evaluated by linearly interpolating its eight neighboring voxels.
- Loss function: The example loss combines an appearance-similarity term based on negative local cross-correlation with a diffusion regularizer penalizing spatial variation in φ.Local means use n = 9, and λ controls regularization strength.
5. Experiments
The experiments evaluate unsupervised atlas-based registration on a large, multi-site brain MRI dataset. Accuracy is assessed through overlap of anatomical segmentations rather than dense ground-truth deformation fields.
- Dataset: The study uses 7829 T1-weighted brain MRI scans from eight publicly available datasets, with preprocessing including affine normalization, brain extraction, cropping, and segmentation.The resulting images are cropped to 160 × 192 × 224.
- Registration task: The atlas-based task registers a fixed external atlas to each randomly selected dataset volume, while all registration is performed in 3D.The atlas serves as the reference volume F and the dataset scan as the moving volume M.
- Evaluation: Dense ground-truth registration is not well-defined, so evaluation uses overlap between corresponding anatomical segmentations across the atlas and warped moving image.The evaluation includes 29 structures containing at least 100 voxels in every test subject.
- Evaluation: Dice score quantifies anatomical volume overlap, with 1 indicating identical structures and 0 indicating no overlap.This metric is used to assess whether the registration field represents accurate anatomical correspondences.
5.3. Baseline Methods
The baseline comparison uses SyN, a state-of-the-art deformable registration algorithm implemented in ANTs with cross-correlation similarity. The experiments use tuned smoothness parameters rather than ANTs defaults.
- Baseline: The study compares VoxelMorph with Symmetric Normalization (SyN), identified as the top-performing algorithm in a prior comparative study.SyN is run through the publicly available ANTs software package with cross-correlation similarity.
- Qualitative comparison: The example figure compares input pairs with warped outputs from VoxelMorph-1 and VoxelMorph-2 using colored anatomical boundaries.The displayed structures are ventricles and hippocampi in coronal MR slices.
- Baseline: The ANTs smoothness parameters were selected through a broad sweep across multiple datasets because the defaults were sub-optimal for these experiments.The tuned parameters are used for the reported baseline experiments.
5.4. Implementation
The implementation uses ADAM-optimized networks, trains separate models across regularization settings, and evaluates held-out performance after preprocessing. VoxelMorph achieves comparable Dice scores to ANTs while testing far faster.
- Optimization: ADAM trains the networks with a 1e-4 learning rate and one volume pair per batch to reduce memory usage.Separate networks are trained for different λ values until convergence.
- Evaluation: Validation Dice selects the network, after which results are reported on a held-out test set.
- Performance: Average Dice scores over all structures and subjects are comparable between VoxelMorph and ANTs.Timing is measured after preprocessing.
5.5. Results
VoxelMorph provides registration accuracy comparable to ANTs and substantially improves on affine alignment. Performance varies by structure and can improve when the model is trained specifically for a subpopulation, while remaining robust across a broad λ range.
- Accuracy: VoxelMorph models achieve comparable Dice measures to ANTs across structures and outperform affine alignment.VoxelMorph-2 performs slightly better than VoxelMorph-1; performance is better on some structures and worse on others relative to ANTs.
- Runtime: 60+ and 150+ times faster on average than ANTs, VoxelMorph-1 and VoxelMorph-2 respectively, using the CPU.On the GPU, both networks compute a registration in under a second.
- Subpopulation training: 1.5% Dice score improvement results from training dataset-specific networks on ABIDE scans rather than all datasets.The evaluation uses unseen ABIDE scans.
- Regularization: λ = 1 for VoxelMorph-1 and λ = 1.5 for VoxelMorph-2 produce the best validation Dice scores.Scores vary slowly across a large λ range, indicating robustness to the regularization choice.
- Regularization: λ = 0 still significantly improves over affine registration despite enforcing no regularization.The authors attribute this likely to shared parameters needing to register all training pairs well.
6. Discussion
The method supports population-specific adaptation and exposes an accuracy–runtime architecture tradeoff. Its speed comes with important scope caveats: Dice is only a proxy, and diffeomorphic properties are not guaranteed.
- Adaptation: Training on ABIDE alone improves test ABIDE Dice scores over training on heterogeneous datasets.The result suggests some model parameters learn properties specific to the training images.
- Architecture tradeoff: VoxelMorph-1 runs significantly faster on CPU and is less than 1 Dice point worse than VoxelMorph-2.Changing network depth and channel count allows application-specific accuracy–runtime tradeoffs.
- Regularization: As λ increases, registration fields become smoother across structural boundaries.The figure visualizes clipped and rescaled x, y, z displacements in RGB channels.
- Limitations: Dice score is a proxy for registration accuracy, while ANTs produces diffeomorphic registrations that the models do not guarantee.Diffeomorphic fields provide invertibility and topology preservation useful in some analyses.
- Practical deployment: A few days of one-time training on a single NVIDIA TITANX GPU precede test-pair registration in less than a second.The approach shifts costly per-pair optimization into dataset-level training.
7. Conclusion
The paper presents an unsupervised learning-based registration approach requiring no ground-truth fields or anatomical landmarks. It reports similar accuracy to state-of-the-art 3D registration while operating orders of magnitude faster.
- The approach requires no supervised information such as ground-truth registration fields or anatomical landmarks.
- On a large-scale, multi-study MR brain dataset, the method achieves similar accuracy to state-of-the-art 3D registration while operating orders of magnitude faster.The paper also reports robustness to regularization and support for population-specific tailoring.