PaperlayerView exact revision on arXiv

Source-linked AI summary

VoxelMorph: A Learning Framework for Deformable Medical Image Registration

Guha Balakrishnan, Amy Zhao, Mert R. Sabuncu, John Guttag, Adrian V. Dalca

arXiv:1809.05231v3cs.CV

TL;DR

Traditional registration optimizes each image pair separately, making large-scale deformable registration slow. VoxelMorph learns a CNN mapping from image pairs to deformation fields, using either intensity-based unsupervised losses or auxiliary segmentations. It achieves state-of-the-art-comparable accuracy with orders-of-magnitude lower runtime, while segmentation-guided training improves test registration accuracy.

  • Problem

    Traditional deformable registration is computationally intensive because it solves a separate optimization problem for each volume pair.

  • Method

    VoxelMorph learns CNN parameters that map image pairs to deformation fields, using intensity-and-smoothness losses or auxiliary anatomical segmentations.

  • Results

    VoxelMorph achieves comparable Dice accuracy to state-of-the-art methods while reducing registration from hours to minutes on CPU and under a second on GPU.

  • Takeaways & Limitations

    VoxelMorph can substantially speed medical image analysis and support segmentation-guided learning for registration when test-time segmentations are unavailable.

Abstract

from arXiv · show

We present VoxelMorph, a fast learning-based framework for deformable, pairwise medical image registration. Traditional registration methods optimize an objective function for each pair of images, which can be time-consuming for large datasets or rich deformation models. In contrast to this approach, and building on recent learning-based methods, we formulate registration as a function that maps an input image pair to a deformation field that aligns these images. We parameterize the function via a convolutional neural network (CNN), and optimize the parameters of the neural network on a set of images. Given a new pair of scans, VoxelMorph rapidly computes a deformation field by directly evaluating the function. In this work, we explore two different training strategies. In the first (unsupervised) setting, we train the model to maximize standard image matching objective functions that are based on the image intensities. In the second setting, we leverage auxiliary segmentations available in the training data. We demonstrate that the unsupervised model's accuracy is comparable to state-of-the-art methods, while operating orders of magnitude faster. We also show that VoxelMorph trained with auxiliary data improves registration accuracy at test time, and evaluate the effect of training set size on registration. Our method promises to speed up medical image analysis and processing pipelines, while facilitating novel directions in learning-based registration and its applications. Our code is freely available at voxelmorph.csail.mit.edu.

I. INTRODUCTION

VoxelMorph replaces costly pairwise registration optimization with a learned CNN-based function that rapidly predicts deformation fields. It supports intensity-based unsupervised training and auxiliary anatomical segmentations, achieving comparable accuracy to state-of-the-art methods with much lower runtime.

  • Traditional deformable registration can require tens of minutes to hours per scan pair, even though GPU implementations reduce this to minutes.
  • VoxelMorph learns a CNN-based registration function from collections of volumes and evaluates it directly on new pairs.This replaces pair-specific test-time optimization with global training of shared function parameters.
  • The framework offers an unsupervised intensity-and-smoothness objective and a second objective that uses anatomical segmentations available during training.
  • On over 3,500 brain scans, unsupervised VoxelMorph achieves accuracy comparable to state-of-the-art registration while requiring less than a minute on CPU and under a second on GPU.The cited baselines require tens of minutes to over two hours on CPU.
  • The study evaluates multiple datasets, registration settings, auxiliary-label configurations, training-set sizes, instance-specific optimization, and hyperparameter sensitivity.

II. BACKGROUND

Medical deformable registration establishes dense nonlinear correspondences between scans despite anatomical variability. VoxelMorph addresses the computational cost of conventional optimization by learning shared registration parameters across volume pairs.

  • Deformable registration enables anatomical comparison across scans with substantial subject-to-subject structural variability.
  • Conventional methods iteratively optimize an energy combining image similarity with deformation regularization for each volume pair.
  • A displacement field commonly represents the transformation as φ = Id + u, where u gives each voxel’s offset.
  • Similarity objectives include intensity error, mutual information, and cross-correlation, while smoothness terms penalize irregular deformation gradients.
  • VoxelMorph replaces pair-specific deformation optimization with dataset-level optimization of shared function parameters.
  • Prior neural registration methods often require ground-truth warp fields, while early unsupervised approaches were limited to subregions, slices, or small transformations.

C. 2D Image Alignment

Learning-based image alignment uses CNNs to infer dense correspondences from image pairs. VoxelMorph extends this approach to n-D volume registration and can incorporate anatomical segmentations during training.

  • Optical flow estimates dense displacement fields for small 2D image motions and is closely related to deformable registration.
  • VoxelMorph extends spatial-transformer-based learning from image alignment to n-D volume registration and auxiliary segmentation-guided training.
  • The method assumes affinely aligned grayscale volumes, with nonlinear misalignment remaining for registration.
  • A CNN maps fixed and moving volumes to a displacement field, with φ = Id + u defining the deformation.
  • During training, a spatial transformation function warps the moving volume so its similarity to the fixed volume can be evaluated.
  • The proposed losses combine image similarity and field smoothness, with an alternative loss incorporating anatomical segmentations.

A. VoxelMorph CNN Architecture

VoxelMorph uses a CNN-based UNet-like function to estimate deformation fields from paired 3D volumes. A differentiable spatial transformation enables training by comparing warped and fixed images.

  • VoxelMorph parameterizes gθ(f, m) with a CNN that estimates a deformation field from two input volumes.The network input is formed by concatenating the moving and fixed volumes.
  • The architecture uses encoder–decoder sections with skip connections to combine coarse contextual features with finer spatial detail.The encoder reduces spatial resolution, while the decoder upsamples and concatenates propagated features for anatomical alignment.
  • The decoder restores the representation to full resolution and further refines it through convolutional layers.The figure specifies several 32-filter convolutions followed by upsampling in the decoder.
  • VoxelMorph computes the warped moving image m ◦φ with a differentiable spatial transformation operation for gradient-based optimization.Voxel locations are shifted by the displacement field and sampled by linear interpolation from neighboring voxels.
  • The training objective can use an unsupervised loss based on the input volumes and the generated registration field.This loss is one of the two proposed training objectives.

1) Unsupervised Loss Function:

The unsupervised objective combines image similarity with deformation smoothness. It supports mean squared error for similar intensity distributions and local cross-correlation for intensity variation across scans.

  • The unsupervised loss combines an appearance-similarity term with a regularizer penalizing local spatial variation in the deformation field.The regularization parameter λ controls the trade-off between these components.
  • Mean squared voxelwise difference is suitable when the fixed and moving images have similar intensity distributions and local contrast.It is one of the two image-similarity functions evaluated for Lsim.
  • Local cross-correlation is more robust to intensity variations across scans and datasets.The method computes local means over an n^3 neighborhood, using n = 9 in the experiments.
  • For cross-correlation, higher CC indicates better alignment, so the similarity loss is defined as −CC(f, m ◦φ).
  • A diffusion regularizer encourages a smooth displacement field by penalizing spatial gradients of displacement u.Spatial gradients are approximated using differences between neighboring voxels.

2) Auxiliary Data Loss Function:

VoxelMorph can incorporate anatomical segmentations available during training through a Dice-based auxiliary loss. This loss is combined with the unsupervised objective because segmentation overlap alone does not enforce image agreement or smoothness.

  • The auxiliary objective uses anatomical segmentation maps to encourage correspondence between matching structures in the fixed and warped moving images.Accurate registration should produce strong overlap for regions representing the same anatomical structure.
  • The method quantifies overlap for each anatomical structure with the Dice score and aggregates the structures into Lseg.A Dice score of 1 denotes perfect anatomical overlap, whereas 0 denotes no overlap.
  • Lseg alone does not encourage smoothness or agreement of image appearance, so it is combined with the unsupervised objective.The combined objective includes a regularization parameter γ.
  • In experiments with affinely aligned images, the combined loss produced significant improvements.The authors note that a multiscale version can be used depending on the quality of the initial alignment.
  • Segmentation channels are transformed separately with linear interpolation so categorical labels can be used in differentiable registration.Each channel is a binary mask for one anatomical structure.

D. Amortized Optimization Interpretation

VoxelMorph replaces pair-specific deformation optimization with global optimization of shared network parameters. The experiments examine brain MRI registration across a large, heterogeneous dataset while recognizing that dense ground-truth fields are not well-defined.

  • D. Amortized Optimization Interpretation: VoxelMorph performs amortized optimization by learning shared parameters θ for a function that estimates registration between image pairs.This replaces costly optimization over φ for each test pair with one global training optimization.
  • Experiments: The method is demonstrated on brain MRI registration using atlas-based experiments and additional analyses of hyperparameters, training-set size, and segmentation data.The atlas serves as a fixed reference volume paired with each dataset volume.
  • Dataset: The evaluation uses a 3731-scan, multi-site, multi-study dataset assembled from publicly available brain MRI datasets.The scans differ in acquisition details, subject age ranges, and health conditions.
  • Qualitative examples: Example MR slices compare input pairs with warped outputs and emphasize whether ventricles, thalami, and hippocampi become anatomically similar.The examples include ventricular expansion or shrinkage and hippocampal stretching.
  • Evaluation: Because dense ground-truth registration is not well-defined for these data, evaluation uses anatomical volume overlap and deformation-field regularity.Multiple fields can yield similarly appearing warped images.

2) Evaluation Metrics:

VoxelMorph is evaluated using Dice overlap, deformation-field regularity, and comparisons with affine alignment, ANTs, and NiftyReg. Across anatomical structures, it achieves comparable accuracy while producing generally smooth, well-behaved deformations.

  • Dice overlap: Dice scores compare affine alignment, ANTs, NiftyReg, and VoxelMorph across anatomical structures and subjects.The visualization averages corresponding left- and right-hemisphere structures and orders structures by average ANTs Dice score.
  • Deformation regularity: Example VoxelMorph (MSE) deformation fields remain smooth within the brain despite substantial differences between moving and fixed images.The warped volume and deformation fields are shown alongside the input images; displacement components are encoded in RGB channels.
  • Dice overlap: VoxelMorph variants perform comparably to ANTs and NiftyReg in average Dice, and outperform affine alignment.VoxelMorph is slightly better on some structures and worse on others, including lateral ventricles and hippocampi.
  • Parameter selection: Validation Dice scores are measured while varying the smoothness regularization parameter λ.Separate networks are trained with different λ values, and the model optimizing validation Dice is selected for testing.
  • Deformation regularity: Deformation regularity is assessed using the number of voxels with non-positive Jacobian determinants.All methods are diffeomorphic at the vast majority of voxels, with 99.4%–99.9% of voxels meeting this criterion.

1) Runtime:

VoxelMorph replaces costly pairwise registration optimization at test time with direct evaluation of a learned function. The paper reports comparable registration quality while highlighting the computational cost of competing CPU-based methods and the lack of GPU implementations for some baselines.

  • Runtime comparison: ANTs requires two or more CPU hours, while NiftyReg requires roughly 39 minutes for the reported setting.These timings follow the shared affine-alignment preprocessing step.
  • Runtime comparison: VoxelMorph computes test-pair registrations by evaluating its learned function instead of optimizing separately for each pair.The framework performs one global function optimization during training and then applies the function to new image pairs.
  • Runtime comparison: VoxelMorph achieves comparable Dice accuracy to ANTs and NiftyReg while operating orders of magnitude faster during testing.The table summarizes Dice, runtime, and non-positive Jacobian counts across affine alignment, ANTs, NiftyReg, and VoxelMorph.
  • Runtime limitation: The reported runtime comparison is limited by the absence of an ANTs GPU implementation and the unavailable NiftyReg GPU version.The authors note that SyN might benefit from GPU implementation, while VoxelMorph’s main advantage is avoiding per-pair optimization.
  • Parameter robustness: Validation Dice varies smoothly across a broad range of λ values, indicating robustness to the regularization parameter.Even λ = 0 improves substantially over affine registration, which the authors attribute to implicit dataset regularization.

D. Training Set Size and Instance-Specific Optimization

The experiments examine how training data size and auxiliary anatomical labels affect registration, including generalization to unseen scans and the trade-off between observed-label accuracy and unobserved-label accuracy.

  • Training set size and optimization: Training-set-size experiments report Dice on the training subset, held-out test set, and test pairs after instance-specific deformation optimization.The study uses MSE and compares amortized VoxelMorph registration with additional per-pair optimization.
  • Auxiliary anatomical labels: Auxiliary-label experiments train models with one, 15, or 30 observed anatomical structures, then evaluate test pairs without using segmentation maps.The observed structures include manually selected single structures or randomly selected subsets of the 30 structures.
  • Observed structures: Auxiliary-data VoxelMorph significantly outperforms unsupervised VoxelMorph and ANTs on observed structures, with the largest p-value below 10^-9 across four settings.Observed-label Dice generally increases as γ increases.
  • Unobserved structures: Auxiliary training also improves Dice for unobserved structures across a range of γ values, despite those structures not being observed during training.When γ becomes large, unobserved-structure accuracy declines as observed-structure accuracy continues improving, consistent with overfitting to observed structures.

2) Training with coarse labels:

The coarse-label experiment merges 30 anatomical structures into four broad groups during training and evaluates registration on the original structures. Moderate auxiliary weighting improves Dice, whereas stronger weighting degrades deformation regularity.

  • Deformation regularity: At γ = ∞, deformation fields are more irregular and warped images are visually less coherent than at γ = 0.01.The auxiliary-only extreme improves observed-structure Dice but removes the smoothness and image-matching terms.
  • Coarse-label setup: The coarse-label setting merges 30 structures into white matter, gray matter, cerebrospinal fluid, and brain stem groups.Accuracy is evaluated on the original, finer-grained structures within these broad groups.
  • Accuracy: A γ value of 0.01 yields mean Dice 0.78±0.03 on FreeSurfer segmentations, roughly 3 Dice points above VoxelMorph without auxiliary information.The improvement has p-value < 10^-10.
  • Deformation regularity: Deformation regularity degrades slowly with increasing γ until the auxiliary-only extreme.Approximately 0.2% of voxels fold at the lowest parameter value, at most 2.3% fold at γ = 0.1, and 10–13% fold at γ = ∞.
  • External evaluation: Testing on manual Buckner40 segmentations shows the same overall behavior, with smaller Dice improvements than on the general test set.The authors suggest this may reflect higher baseline Dice scores on Buckner40.

VI. DISCUSSION AND CONCLUSION

VoxelMorph matches established registration accuracy while substantially reducing runtime, and auxiliary anatomical information can further improve test-time registration. Its amortized learning framework remains effective with modest training sets and extends beyond the demonstrated brain-scan setting.

  • VoxelMorph matches ANTs and NiftyReg in Dice score while reducing computation from hours to minutes on CPU and under a second on GPU.
  • Auxiliary anatomical segmentations can improve test registration accuracy while preserving smooth deformations and the runtime advantage.
  • With 100 training images, VoxelMorph achieves state-of-the-art registration quality on both training and test sets; instance-specific optimization adds one Dice point.
  • The benefit of auxiliary data depends on segmentation quality and quantity: partial or coarse labels can improve registration across test labels.
  • VoxelMorph is not limited to one anatomy or image type and may support applications such as cardiac MR, lung CT, and multimodal registration.
Loading 1809.05231v3…