Source-linked AI summary
Unsupervised Learning for Fast Probabilistic Diffeomorphic Registration
Adrian V. Dalca, Guha Balakrishnan, John Guttag, Mert R. Sabuncu
TL;DR
Existing registration methods trade computational cost against deformation flexibility, supervision, topology guarantees, or uncertainty estimation. The paper develops an unsupervised probabilistic CNN-based inference framework with differentiable diffeomorphic integration, achieving fast, accurate registration with uncertainty estimates and diffeomorphic guarantees.
Problem
Existing methods can be computationally intensive, require supervised deformation labels, use restricted deformation models, lack diffeomorphic guarantees, or provide no probabilistic uncertainty estimates.
Method
The paper uses variational inference in a probabilistic generative model with a CNN, spatial transform layers, and differentiable scaling-and-squaring layers for unsupervised diffeomorphic registration.
Results
The algorithm achieves state-of-the-art registration accuracy and fast runtimes while producing diffeomorphic fields and registration uncertainty estimates.
Takeaways & Limitations
The method infers new image-pair registrations in under a second while combining speed with diffeomorphic guarantees and natural uncertainty estimates.
Takeaways & Limitations
The model uses a stationary velocity field representation and encourages spatial smoothness through a graph-Laplacian prior covariance structure.
Abstract
from arXiv · showhide
Traditional deformable registration techniques achieve impressive results and offer a rigorous theoretical treatment, but are computationally intensive since they solve an optimization problem for each image pair. Recently, learning-based methods have facilitated fast registration by learning spatial deformation functions. However, these approaches use restricted deformation models, require supervised labels, or do not guarantee a diffeomorphic (topology-preserving) registration. Furthermore, learning-based registration tools have not been derived from a probabilistic framework that can offer uncertainty estimates. In this paper, we present a probabilistic generative model and derive an unsupervised learning-based inference algorithm that makes use of recent developments in convolutional neural networks (CNNs). We demonstrate our method on a 3D brain registration task, and provide an empirical analysis of the algorithm. Our approach results in state of the art accuracy and very fast runtimes, while providing diffeomorphic guarantees and uncertainty estimates. Our implementation is available online at http://voxelmorph.csail.mit.edu .
1 Introduction
The paper targets fast, unsupervised registration that retains diffeomorphic topology preservation and probabilistic uncertainty, addressing limitations of conventional and recent learning-based methods. It formulates registration as variational inference and uses differentiable scaling-and-squaring components to learn diffeomorphic transformations.
- Traditional diffeomorphic registration preserves topology but often requires substantial time and computational resources for each image pair.
- Learning-based registration methods can require ground-truth deformation fields derived from conventional tools, introducing bias and preprocessing demands.
- Earlier unsupervised methods were demonstrated with constrained deformation models and limited 2D or 3D inputs, while another approach lacked topology guarantees and uncertainty estimates.
- The paper formulates registration as variational inference in a probabilistic generative model and derives a CNN-based learning algorithm with an intuitive cost function.
- Novel diffeomorphic integration layers combined with transform layers enable unsupervised end-to-end learning of diffeomorphic registration.
- 1.1 Diffeomorphic Registration: Stationary velocity fields are integrated numerically through scaling and squaring, producing the final diffeomorphic transformation by recursive composition.
2 Methods
The method frames 3D deformable registration as unsupervised variational inference, using a CNN to estimate a probabilistic velocity field and differentiable layers to produce diffeomorphic warps with uncertainty estimates.
- Generative model: A generative model represents the moving image as y ◦φz, where a latent variable z parameterizes the transformation and posterior registration probability is estimated.The model treats x as a noisy observation of the warped image and supports MAP registration and uncertainty estimation.
- Generative model: The prior uses a multivariate normal distribution over z, with precision Λz = λL encouraging spatial smoothness of the stationary velocity field.L is the Laplacian of a voxel-neighborhood graph, and λ controls the velocity-field scale.
- Variational learning: Variational inference learns an approximate posterior qψ(z|x; y) modeled as a multivariate normal with diagonal covariance, using a convolutional network to predict its mean and variance.The network parameters are optimized with stochastic gradient methods on the variational lower-bound objective.
- Variational learning: The loss combines image similarity between x and y ◦φzk with a KL term that keeps the posterior near the prior and spatially smooths the posterior mean.The experiments use K = 1 sampled latent velocity field; σ^2 and λ are fixed hyper-parameters.
- Diffeomorphic transformation: Differentiable scaling-and-squaring layers exponentiate sampled stationary velocity fields into diffeomorphic deformation fields, which a spatial-transform layer applies to y.The network outputs posterior parameters and the warped image; the implementation uses T = 7 scaling-and-squaring steps.
- Registration and uncertainty: For new image pairs, the network predicts the MAP velocity as the posterior mean, computes its deformation field, and estimates uncertainty from posterior covariance and sampled deformations.Multiple samples are propagated through the diffeomorphic layers to obtain empirical deformation-field covariance.
3 Experiments
The experiments evaluate 3D atlas-based brain registration on a large multi-site MRI dataset, comparing accuracy, runtime, diffeomorphic properties, and uncertainty. The method matches baseline Dice scores, runs much faster, produces nearly topology-preserving fields, and provides uncertainty estimates.
- Data and task: The study focuses on 3D atlas-based registration using 7829 T1-weighted brain MRI scans from eight public datasets.The data are split into 7329 training, 250 validation, and 250 test volumes, with training performed unsupervised.
- Evaluation: Registration quality is evaluated by propagating anatomical segmentations to an atlas and measuring volume overlap with the Dice metric.The evaluation also counts voxels with non-positive Jacobian determinants to assess diffeomorphic behavior.
- Accuracy: All methods achieve comparable Dice scores across anatomical structures and overall.Figure 3 summarizes scores for structures including the brain stem, thalamus, ventricles, hippocampus, and cerebral cortex.
- Runtime and topology: Our method and VoxelMorph register two images in less than a second on a GPU, while our method takes less than a minute on a CPU.These runtimes exclude preprocessing common to all methods; ANTs has no GPU implementation reported here.
- Runtime and topology: Our method produces positive Jacobians at nearly all voxels, whereas baseline flow fields contain a few thousand non-positive-Jacobian locations.Increasing spatial regularization can alleviate baseline violations but lowers Dice performance.
- Uncertainty: Velocity-field uncertainty is lower near anatomical structure edges and higher in homogeneous regions such as white matter or ventricle interiors.This edge-related pattern is less obvious in the final registration-field uncertainty.
4 Conclusion
The paper proposes probabilistic, unsupervised learning for diffeomorphic image registration using a CNN and differentiable scaling-and-squaring layers. It reports fast inference while retaining diffeomorphic guarantees and providing uncertainty estimates.
- Conclusion: The paper proposes a probabilistic model for diffeomorphic image registration and derives an associated learning algorithm.The algorithm uses a convolutional neural network and an intuitive loss function.
- Conclusion: Novel differentiable scaling-and-squaring layers enable unsupervised, end-to-end learning of diffeomorphic registrations.The derivation can represent a deformation with a low-dimensional embedding or the displacement field itself.
- Conclusion: The algorithm infers registration for new image pairs in under a second while offering diffeomorphic guarantees and natural uncertainty estimates.The paper contrasts this speed with traditional methods and these guarantees with recent learning-based methods.