VoxelMorph: A Learning Framework for Deformable Medical Image Registration (1809.05231v3)

Abstract: 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.

• The paper introduces a CNN-based framework that generates deformation fields rapidly, outperforming traditional registration methods in speed and efficiency.
• It achieves state-of-the-art Dice scores through unsupervised learning and auxiliary segmentation, ensuring high registration accuracy.
• The approach’s amortized optimization minimizes per-image computation, making it scalable for large-scale medical image analysis.

VoxelMorph: An Overview of a Learning Framework for Deformable Medical Image Registration

The paper "VoxelMorph: A Learning Framework for Deformable Medical Image Registration" introduces an efficient and scalable neural network-based method for performing deformable registration of medical images. The authors, Guha Balakrishnan et al., leverage convolutional neural networks (CNNs) to circumvent the computational burdens typically associated with traditional registration methods, providing a framework that is both rapid and accurate.

Main Contributions

VoxelMorph proposes a learnable registration function that maps an input image pair to a deformation field that aligns them. This is realized by parameterizing the registration function via a CNN, facilitating the rapid computation of deformation fields during testing. The framework adopts two training strategies: an unsupervised mode based on intensity-based matching objectives, and a mode leveraging auxiliary segmentation data available during training. Several critical contributions and findings from the paper include:

• Rapid Registration: VoxelMorph significantly reduces the computational time for deformable registration. Traditional methods can take tens of minutes to hours per image pair, whereas VoxelMorph operates in less than a minute on a CPU and under a second on a GPU.
• Accuracy of Unsupervised Models: The unsupervised VoxelMorph models achieve registration accuracy comparable to state-of-the-art methods measured by Dice scores.
• Auxiliary Data Utilization: Utilizing auxiliary segmentation data during training enhances registration accuracy. This is particularly beneficial when segmentations for some structures or coarse labels are available during the training phase.
• Amortized Optimization: The framework performs amortized optimization—optimizing a global function during training rather than solving an optimization problem for each image pair, thus offering implicit regularization.

Methodology

The VoxelMorph framework uses a CNN to model the registration function, which outputs a dense displacement field that aligns the input volumes. The CNN architecture is inspired by the UNet model, comprising an encoder-decoder structure with skip connections to handle various scales of image features.

Loss Functions

Two loss functions are introduced:

1. Unsupervised Loss (\mathcal{L}_{us}): Combines image similarity (using mean squared error or local cross-correlation) and smoothness regularization.
2. Auxiliary Loss (\mathcal{L}_{a}): Integrates unsupervised loss with an additional segmentation loss based on Dice scores, leveraging auxiliary segmentation data when available.

Experimental Results

The paper presents thorough experiments involving a multi-paper dataset with 3731 T1-weighted brain MRI scans. The results demonstrate that VoxelMorph can:

• Achieve Dice scores comparable to traditional methods like ANTs (SyN) and NiftyReg.
• Provide rapid registration—over 150 times faster than conventional methods on a CPU.
• Handle diverse anatomical variations and provide smooth deformation fields, crucial for medical image analysis.

Implications and Future Research

VoxelMorph's approach has significant implications for medical image processing pipelines. It offers a way to streamline the registration process, making it feasible to handle large datasets efficiently—an essential capability for population studies and large-scale medical analyses.

The practical benefits include:

• Scalability: Ability to process medical scans quickly, facilitating large-scale analyses.
• Flexibility: Application across various medical imaging tasks, irrespective of the specific anatomy or imaging modality.
• Enhanced Accuracy: Potential for improved registration accuracy using auxiliary data without the need for ground truth deformable fields.

Future research could explore extending VoxelMorph's approach to multi-modal registrations, incorporating diffeomorphic constraints to ensure topological regularity and leveraging more advanced network architectures potentially augmented by attention mechanisms or recurrent connections for iterative refinement.

Conclusion

VoxelMorph represents a significant advancement in the field of medical image registration, providing a powerful, efficient, and flexible framework for deformable image alignment. Its capacity to integrate auxiliary data and perform rapid, accurate registrations sets a new standard for the development of neural network-based medical image registration tools. The approach demonstrates that global function optimization can nearly match instance-specific optimization, promising smoother, more generalizable deformable registrations and paving the way for wide-ranging applications in medical imaging.