PaperlayerView exact revision on arXiv

Source-linked AI summary

Unsupervised Learning of Probabilistic Diffeomorphic Registration for Images and Surfaces

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

arXiv:1903.03545v2cs.CVcs.GR

TL;DR

The paper addresses the gap between computationally intensive classical registration and fast learning-based methods that may lack diffeomorphic guarantees or probabilistic uncertainty estimates. It derives an unsupervised variational-inference algorithm using a probabilistic generative model and CNNs, achieving fast, accurate image and surface registration with diffeomorphic guarantees. The framework also incorporates anatomical surfaces during training and can estimate deformation uncertainty.

  • Problem

    Classical registration is computationally intensive, while learning-based methods may lack unrestricted deformation models, supervision-free training, diffeomorphic guarantees, or probabilistic uncertainty estimates.

  • Method

    The method performs unsupervised variational inference in a probabilistic generative model, using CNNs, differentiable diffeomorphic integration layers, and spatial transformation operations.

  • Results

    The algorithm achieves state-of-the-art registration accuracy with fast runtime and diffeomorphic deformations, including registration of new image pairs in under a second.

  • Takeaways & Limitations

    The framework connects classical registration guarantees with learning-based speed, supports anatomical surfaces during training, and provides estimates of registration uncertainty.

  • Takeaways & Limitations

    Surface-only registration is beyond the paper’s scope, and the surface loss requiring corresponding points cannot be computed directly because independently provided segmentations lack those correspondences.

Abstract

from arXiv · show

Classical 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 build a connection between classical and learning-based methods. We present a probabilistic generative model and derive an unsupervised learning-based inference algorithm that uses insights from classical registration methods and makes use of recent developments in convolutional neural networks (CNNs). We demonstrate our method on a 3D brain registration task for both images and anatomical surfaces, and provide extensive empirical analyses. Our principled approach results in state of the art accuracy and very fast runtimes, while providing diffeomorphic guarantees. Our implementation is available at http://voxelmorph.csail.mit.edu.

1 Introduction

The paper connects classical diffeomorphic registration with fast learning-based inference through a probabilistic, unsupervised framework. It extends the framework to anatomical surfaces and reports fast, accurate registration while preserving classical deformation guarantees.

  • 1 Introduction: Classical registration offers rigorous deformation theory but requires computationally intensive optimization for each image pair, whereas learning-based methods are faster but often lack topology-preserving guarantees.The paper positions its approach as combining strengths from both paradigms.
  • 1 Introduction: The resulting framework registers new image pairs in under a second on a GPU while providing diffeomorphic guarantees and state-of-the-art registration accuracy.The paper implements the method in VoxelMorph and evaluates image and surface registration extensions.
  • 1 Introduction: The method formulates registration as variational inference in a probabilistic generative model and learns a global convolutional neural network through unsupervised end-to-end training.The formulation is designed to output distributions of deformations rather than only point estimates.
  • 1 Introduction: The framework extends probabilistic registration to anatomical surfaces, allowing segmentations during training while not requiring surfaces at test time.The surface model combines surface and image-based training within a single framework, using a differentiable surface distance function.
  • 1 Introduction: Diffeomorphic transforms are generated by integrating stationary velocity fields, with scaling and squaring implemented as differentiable network operations.The deformation is obtained by integrating the velocity field from the identity transformation to the final registration field.

2 Methods

The method frames registration as unsupervised variational inference over a probabilistic deformation model, using a CNN to predict velocity-field distributions. Differentiable integration and spatial transformation layers produce diffeomorphic warps, while optional surface losses support anatomical surface alignment.

  • Generative Model: A probabilistic generative model represents the fixed image as a noisy observation of the moving image warped by a latent transformation.The latent variable is a stationary velocity field that parametrizes a diffeomorphism, with a smoothness-inducing prior.
  • Learning: Variational inference uses a CNN to estimate the approximate posterior mean and diagonal voxel-wise variance of the velocity field without ground-truth registrations.The learned parameters optimize a variational lower bound using stochastic gradient methods.
  • Neural Network Framework: The network samples a velocity field, integrates it through differentiable scaling-and-squaring layers, and spatially transforms the moving image to produce a warped image.The architecture is end-to-end differentiable, enabling stochastic-gradient training.
  • Learning: The variational loss encourages similarity between the fixed image and warped moving image, posterior proximity to the prior, and spatial smoothness of the posterior mean.The implementation uses one sample, K = 1, for the expectation approximation.
  • Registration: At test time, the most likely velocity field is predicted and integrated into a deformation in less than a second on a GPU.The diagonal covariance is not used for the test-time deformation but enables estimation of deformation uncertainty; uncertainty analysis is left for future work.

3 Method Extensions

The framework extends probabilistic diffeomorphic registration to anatomical surfaces and models uncertainty with smoother covariance approximations. Surface losses use distance transforms because independently segmented surfaces lack point correspondences.

  • 3.1 Surface-based Registration: The model extends image registration to anatomical surfaces, enabling segmentation information during training within the same probabilistic framework.
  • 3.1 Surface-based Registration: Surface coordinates are modeled by warping moving-surface points with the diffeomorphism, but independently supplied segmentations do not provide corresponding points.
  • 3.1 Surface-based Registration: Distance transforms approximate the surface-matching term by measuring distances in both directions between warped and fixed surfaces.The implementation precomputes the fixed surface distance transform, samples 100,000 moving-surface points, and evaluates warped locations.
  • 3.1 Surface-based Registration: At test time, images alone are required for inference, while training can evaluate both warped images and warped surfaces.
  • 3.1 Surface-based Registration: Surface-only registration removes image likelihood terms and uses segmentation maps as network inputs, but remains beyond this paper’s scope.
  • 3.2 Non-diagonal Covariance: A smoothed non-diagonal covariance models spatially correlated velocity uncertainty, producing smoother training fields; both covariance approximations remain accurate and diffeomorphic.The smoothing scale is matched to the prior, and its effect diminishes as λ increases.

4 Experiments

Experiments evaluate accuracy, runtime, and diffeomorphic behavior across image-registration methods. VoxelMorph-diff combines comparable accuracy with much faster runtime and nearly no folding voxels.

  • The experiments assess registration accuracy, runtime, diffeomorphic behavior, surface extensions, and the effects of integration layers.
  • VoxelMorph-diff achieves comparable Dice scores and orders-of-magnitude faster runtime than classical baselines while maintaining nearly zero non-positive Jacobian locations.
  • Across anatomical structures, ANTs, NiftiReg, VoxelMorph (CC), and VoxelMorph-diff show generally comparable Dice results, with each method better on some structures.
  • The evaluation focuses on atlas-based registration, referring to the proposed method as VoxelMorph-diff within the VoxelMorph framework.

4.1 Experiment setup

The experiments use a large, multisite brain-MRI dataset with anatomical segmentations and evaluate atlas registration using overlap, surface distance, and folding measures. Comparisons include classical, unsupervised CNN, and supervised baselines.

  • 4.1.1 Data and Preprocessing: The dataset contains 3,731 T1-weighted brain MRI scans from eight publicly available datasets, with standardized preprocessing and unsupervised training.
  • 4.1.1 Data and Preprocessing: Figures show representative warped MR slices and deformation fields, while surface results report maximum and median Euclidean distance and mean Dice.
  • 4.1 Experiment setup: Registration quality is measured by propagating atlas segmentations and computing Dice overlap; surface experiments additionally use Euclidean surface distance.
  • 4.1 Experiment setup: Diffeomorphic behavior is evaluated by counting voxels with non-positive Jacobian determinants, which indicate local folding.
  • 4.1 Experiment setup: The study compares VoxelMorph-diff with tuned ANTs SyN and the non-diffeomorphic CNN baseline VoxelMorph.
  • 4.1 Experiment setup: A supervised VoxelMorph-diff baseline is trained on deformation fields generated by registering more than 650 atlas–MRI training pairs with NiftyReg.

4.2 Image Registration

VoxelMorph-diff achieves state-of-the-art Dice accuracy and fast runtimes while producing nearly folding-free diffeomorphic registration fields. Learning-based methods register images in under a minute and a half on CPU and under a second on GPU.

  • Accuracy and runtime: VoxelMorph-diff achieves state-of-the-art Dice results and runtimes while producing diffeomorphic registration fields with nearly no folding voxels per scan.The method operates within a probabilistic framework.
  • Comparison: All methods have comparable Dice results across structures and overall, except the supervised method, which has a slight Dice loss despite more diffeomorphic deformations.The supervised network was trained on 650 subjects.
  • Accuracy and runtime: Learning-based methods register two images in less than a second on a GPU and less than a minute and a half on a CPU.Reported runtimes exclude preprocessing common to all methods.
  • Deformation regularity: VoxelMorph-diff produces zero folding voxels for most scans, averaging less than one folding voxel per test scan.Baseline methods contain a few thousand locations of non-positive Jacobians per scan.

4.3 Image and Surface Registration

VoxelMorph-surf incorporates anatomical surfaces during training alongside images and improves registration near the observed structure while preserving low folding counts. Its gains are significant across surface distance and Dice measures for most desired structures.

  • Model and evaluation: VoxelMorph-surf uses anatomical segmentation alongside images during training to evaluate a joint generative surface model.Separate networks are learned for five anatomical structures.
  • Registration accuracy: VoxelMorph-surf significantly improves Euclidean surface distance and Dice over image-only VoxelMorph-diff for most desired structures.The comparison is performed on the test set.
  • Deformation regularity: The joint model increases folding voxels only slightly to an average 3.5 voxels per volume despite improved accuracy in both metrics.This remains orders of magnitude fewer than the baseline methods.
  • Registration accuracy: Approximately 9 Dice points of hippocampus improvement and more than three voxels of maximum surface-distance improvement preserve fewer than three folding voxels per scan.These gains are reported when hippocampus surfaces are available during training.

4.4 Analysis

Analyses show that scaling and squaring efficiently produces accurate, regular, and invertible deformations, while variational covariance approximations yield similar results at reasonable regularization. Controlled experiments confirm accurate, folding-free, invertible warps.

  • Regularization analysis: Dice scores vary smoothly across λ values from 0.5 to 100, with λ = 20 used in the reported experiments.The analysis uses 50 validation scans and σ²_I = 0.02.
  • Controlled experiments: Disk-to-C-shape deformation fields achieve accurate registration with no folding voxels, and composing forward and inverse transforms returns the grid to identity.The experiments vary C-shape radius and thickness.
  • Integration analysis: After four scaling and squaring steps, VoxelMorph-diff reaches maximum Dice, while additional steps sharply reduce folding voxels.Runtime increases modestly and remains much smaller than the rest of the network computation.
  • Integration analysis: Five scaling and squaring operations require 0.06 ± 0.01 seconds, compared with 0.53±0.01 seconds for equivalent quadrature integration.Quadrature requires 64 operations and prohibitive memory.
  • Integration analysis: Increasing scaling and squaring layers improves deformation-field properties at negligible runtime cost, motivating T = 7 test-time squaring steps.The analysis characterizes accuracy, runtime, regularity, and invertability.
  • Variational covariance: For a given λ, diagonal and non-diagonal covariance approximations produce comparable accuracy and folding counts after integration.The diagonal approximation can add voxel-independent noise, but smoothing and diffeomorphic integration yield well-behaved fields at reasonable λ values.

5 Discussion and Conclusion

The paper connects classical probabilistic registration with unsupervised CNN inference by integrating stationary velocity fields into diffeomorphic transforms. It extends the framework to anatomical surfaces, achieving under-second inference while retaining diffeomorphic guarantees.

  • Contributions: The method uses variational inference on a probabilistic model with a CNN and differentiable scaling-and-squaring integration for unsupervised diffeomorphic registration.The formulation connects classical registration methods with learning-based approaches.
  • Surface extension: The surface model incorporates segmentations during training within the joint framework without requiring surfaces at test time.This provides an anatomical surface deformation model alongside image registration.
  • Conclusions: New image-pair registration takes under a second, while the method offers diffeomorphic guarantees compared with recent learning-based methods.The surface extension improves registration while preserving low runtime and diffeomorphisms.
  • Related implication: With very limited training data, VoxelMorph deformations can initialize a classical method and enable faster convergence.This conclusion is attributed to a related prior result.
  • Extensions: The variational covariance can estimate deformation uncertainty at each voxel, and the formulation is described as generalizable to other deformation representations and applications.Suggested applications include biomedical segmentation, population analysis, and multiple-surface registration.

Supplementary Material

The supplementary material derives the variational lower bound and final loss for the generative surface model, using conditional-independence assumptions and sampled expectations.

  • The expectation is approximated with samples of the latent velocity field, producing the final loss after combining the additional surface term with the preceding loss.
  • The derivation begins by minimizing the KL divergence between the true and approximate posteriors.
  • The factorization separates the fixed-image likelihood from the fixed-surface likelihood given the deformation.
  • The objective expands into a posterior-prior KL term and expected negative log-likelihood terms for the fixed image and fixed surface.
  • This factorization assumes conditional independence between the fixed image and anatomical surfaces given the moving image and deformation, and between the fixed surface and either image given the moving surface and deformation.

Overview figure with surface loss

The end-to-end unsupervised architecture infers a probabilistic velocity field, converts it into a diffeomorphic deformation, and warps the moving image.

  • The network defψ(m, f) takes moving and fixed images as input and predicts posterior mean and variance parameters for the velocity field.
  • A velocity field is sampled from the approximate posterior and transformed into a diffeomorphic deformation field through differentiable squaring and scaling layers.
  • The resulting deformation field spatially transforms the moving image to produce the warped image m ◦φz.

Additional Figures

The additional figures show image- and surface-driven registration examples across multiple subjects, comparing the proposed methods with ANTs and visualizing anatomical boundaries and deformation fields.

  • Figure 13 compares moving images, atlases, and warped images from VoxelMorph-diff and ANTs across rows representing different scans.
  • The Figure 13 examples overlay boundaries of the ventricles, thalami, and hippocampi, while warped grids and RGB images visualize the registration field.
  • VoxelMorph and NiftyReg examples are omitted from Figure 13 because they are visually similar to the ANTs results.
  • Figure 14 presents example surface-driven results for three subjects, comparing VoxelMorph-diff, VoxelMorph-surf, and ANTs using cropped MR slices and ventricle boundaries.
  • The red-box examples in Figure 14 highlight improved ventricle segmentation from the top row to the bottom row.
Loading 1903.03545v2…