Source-linked AI summary
Learning a Probabilistic Model for Diffeomorphic Registration
Julian Krebs, Hervé Delingette, Boris Mailhé, Nicholas Ayache, Tommaso Mansi
TL;DR
The paper addresses how to learn a compact probabilistic model of deformations for registration and deformation analysis. It uses an unsupervised CVAE with diffeomorphic, symmetric, and multi-scale deformation modeling. On cardiac cine-MRI, it reports state-of-the-art registration performance alongside deformation transport and disease clustering capabilities.
Problem
The paper seeks a learned low-dimensional probabilistic deformation model that supports registration, comparison, generation, and transport of deformations.
Method
An unsupervised CVAE learns deformation encodings, while stationary velocity-field exponentiation, symmetric losses, spatial regularization, and multi-scale estimation constrain registration.
Results
The multi-scale model achieved statistically significant registration improvements over comparable methods, produced smoother deformation fields, and supported deformation transport and disease clustering.
Takeaways & Limitations
The learned latent space provides a shared representation for generating, comparing, transporting, and analyzing deformations across subjects.
Takeaways & Limitations
The authors question whether a multivariate Gaussian is the appropriate latent-space prior and leave alternative distributions for future work.
Abstract
from arXiv · showhide
We propose to learn a low-dimensional probabilistic deformation model from data which can be used for registration and the analysis of deformations. The latent variable model maps similar deformations close to each other in an encoding space. It enables to compare deformations, generate normal or pathological deformations for any new image or to transport deformations from one image pair to any other image. Our unsupervised method is based on variational inference. In particular, we use a conditional variational autoencoder (CVAE) network and constrain transformations to be symmetric and diffeomorphic by applying a differentiable exponentiation layer with a symmetric loss function. We also present a formulation that includes spatial regularization such as diffusion-based filters. Additionally, our framework provides multi-scale velocity field estimations. We evaluated our method on 3-D intra-subject registration using 334 cardiac cine-MRIs. On this dataset, our method showed state-of-the-art performance with a mean DICE score of 81.2% and a mean Hausdorff distance of 7.3mm using 32 latent dimensions compared to three state-of-the-art methods while also demonstrating more regular deformation fields. The average time per registration was 0.32s. Besides, we visualized the learned latent space and show that the encoded deformations can be used to transport deformations and to cluster diseases with a classification accuracy of 83% after applying a linear projection.
I. INTRODUCTION
The paper motivates deformation analysis beyond pairwise registration and proposes learning a deformation model directly from training images. This supports comparing, encoding, and transporting deformations while addressing the optimization and registration challenges of prior approaches.
- Motivation: Deformable registration finds voxel correspondences between image pairs for medical analysis, including tracking structures and disease progression.The resulting deformation fields can support preoperative/postoperative studies, multimodal matching, and cardiovascular diagnosis or therapy selection.
- Contribution: The proposed algorithm learns a deformation model directly from training images and encodes similar deformations close together for comparison and transport.The low-dimensional encoding also supports synthetic deformation generation for individual images.
- Motivation: Traditional registration is an ill-posed optimization problem requiring similarity metrics, predefined deformation models, and regularization based on prior assumptions.Proposed regularizers include elastic and diffusion-based energies, while diffeomorphic transforms are folding-free and invertible.
- Related approaches: Learning-based registration seeks to reduce lengthy iterative optimization, but supervised methods generally require proxy mappings from existing algorithms, simulations, or both.Unsupervised methods instead optimize image similarity, often with penalization or smoothing terms.
- Deformation analysis: Existing deformation-analysis methods compare or transport deformations, but transport typically requires multiple registrations, including difficult inter-subject registrations.Parallel transport methods use SVF or LDDMM parameterizations to apply one subject’s temporal deformation to another.
C. Learning-based Generative Latent Variable Models
The paper positions its method as a learned global probabilistic latent space for diffeomorphic registration. Unlike prior probabilistic representations, it supports sampling, deformation transport, and analysis across image pairs without relying on a unique template.
- Research gap: Prior probabilistic deformation representations had not been learned from large datasets or used for sampling new deformations.The paper identifies the absence of global latent models that organize similar deformations in a probabilistic subspace.
- Generative latent model: The method learns an unsupervised low-dimensional global latent space in an encoder-decoder network, placing similar deformations close together.The space can generate infinitely many deformations for any single image from the learned data distribution.
- Generative latent model: Diffeomorphic deformations are reconstructed by decoding a latent code under the constraint of a moving image.Applying a code from one image pair to a new constraining image enables deformation transport and sampling.
- Generative latent model: A CVAE reconstructs the fixed image by warping the moving image, conditioning the decoder on the moving image to favor deformation rather than appearance encoding.This implicit decoupling supports transporting a latent deformation code to another case.
- Evaluation: The framework is evaluated on cardiac MRI registration and on latent-code structure and deformation transport.The paper also extends preliminary work with probabilistic derivations, multi-scale estimation, latent-space analysis, and transport evaluation.
II. METHODS
The method models registration with a low-dimensional latent deformation code that generates warped moving images. It uses stationary velocity fields, diffeomorphic exponentiation, multi-scale decoding, and a local cross-correlation likelihood.
- Registration formulation: Registration seeks a transformation T_z that warps moving image M to match fixed image F while balancing similarity D and spatial regularization R.Recent unsupervised methods learn the similarity objective using stochastic gradient descent and a spatial transformer layer.
- Probabilistic model: The proposed model learns a probabilistic low-dimensional deformation vector z from example image pairs and decodes it by warping the moving image.Sampling z from its prior generates new deformations based on the learned model.
- Multi-scale diffeomorphic model: Deformation fields are parameterized by stationary velocity fields v, with transformations φ = exp(v) represented across multiple image scales.Each scale defines a decoding function conditioned on z and the corresponding moving image.
- Training objective: At training time, decoded velocities are exponentiated and used to warp moving images, with likelihoods combined across scales.The likelihood uses local cross-correlation rather than the commonly used Gaussian or SSD criterion.
- Multi-scale diffeomorphic model: The generative process models the fixed-image likelihood pθ(F|z, M), while regularized registration uses regularized velocities v∗ in pθ(F|v∗, M).The regularized formulation links regularized and ordinary velocities through a Gaussian distribution.
- Training objective: Symmetric warping uses M∗s = M^s ◦ exp(v_s/2) and F∗s = F^s ◦ exp(−v_s/2) within the multi-scale formulation.This symmetric construction is part of the image-warping objective described for each scale.
1) Learning the constrained deformation encoding:
The constrained deformation-encoding stage addresses how to represent latent variables and how to evaluate the model when integrating over latent codes is intractable. Variational inference replaces the exact posterior with a learned encoder distribution.
- Latent representation: The method must determine what information latent variables z represent while avoiding a direct interpretation of individual latent dimensions.VAEs instead require samples from z to follow a simple prior distribution p(z).
- Variational inference: The integral over z is intractable, so the model samples latent codes likely to have produced the observed fixed and moving images.This requires approximating the posterior p(z|F, M) rather than evaluating it exactly.
- Variational inference: The encoder learns q_ω(z|F, M) as an approximation to the intractable posterior over latent deformations.The encoder and decoder are optimized jointly through the variational objective.
- Variational inference: Maximizing the evidence lower bound replaces direct likelihood maximization with a differentiable objective containing the encoder and decoder.The KL divergence measures how closely q_ω approximates the true posterior and ideally vanishes with sufficient encoder capacity.
2) Optimizing the ELBO:
The ELBO combines a KL divergence aligning the encoder distribution with a Gaussian prior and a reconstruction term estimated with stochastic sampling. Multi-scale likelihoods combine additively across scales.
- ELBO terms: The objective jointly optimizes the KL divergence between qω(z|F, M) and the prior p(z), plus the expected reconstruction log-likelihood.The encoder distribution is Gaussian with diagonal covariance, and its KL term has a closed-form expression.
- ELBO terms: The KL expression uses k for the dimensionality of the distribution.
- Stochastic optimization: One latent sample estimates the reconstruction expectation, reducing computation while training uses stochastic gradient descent over sampled image pairs.
- Stochastic optimization: The reparametrization trick enables back-propagation through latent sampling.
- Multi-scale objective: The product of likelihoods across scales becomes a sum of log-likelihoods.
B. Introducing regularization on velocities
The framework regularizes neural-network velocity fields spatially before decoding, using a probabilistic formulation that yields smoothed multi-scale velocities and supports diffusion-like Gaussian filtering.
- Velocity regularization: Velocity fields generated by neural networks are spatially regularized with Gaussian convolution to control their smoothness.The Gaussian has standard deviation σG, and alternative quadratic Tikhonov or prior-enforcing functionals are possible.
- Probabilistic formulation: The probabilistic formulation treats neural-network velocities as random variables and introduces v∗ for their regularized counterparts.The posterior over v is modeled as concentrated at the network output, while v∗ is linked through a Gaussian distribution.
- Probabilistic formulation: A diffusion-like prior is imposed on v∗, and the regularized velocity estimate is obtained by maximum a posteriori inference.
- Gaussian smoothing: The regularized velocity ˆv is equivalent to Gaussian convolution, ˆv = GσG ∗v, through the Heat equation.
- Multi-scale decoding: The decoder operates on regularized velocities, which are combined across scales within a multi-scale loss.The resulting loss is defined per training image pair and sample.
C. Network architecture
The network uses a conditional variational autoencoder for multi-scale registration, producing latent codes, velocity fields, diffeomorphisms, and warped images conditioned on the moving image.
- Network architecture: The encoder-decoder receives moving and fixed images and outputs z, velocities, deformation fields, and warped moving images at multiple scales.Strided convolutions form the encoder, while fully connected bottleneck layers produce μ, σ, and z.
- Network architecture: The decoder conditions its deconvolution layers on sub-sampled moving-image information.This conditioning is implemented by concatenating each layer’s output with sub-sampled versions of M.
- Diffeomorphic warping: Stationary velocity fields are exponentiated by a scaling-and-squaring layer to obtain diffeomorphisms, which a spatial transformer uses to warp images.
- Network architecture: The probabilistic multi-scale network maps deformations to latent variables and decodes velocities and diffeomorphisms while conditioned on M.The latent dimensionality is illustrated with d = 32.
- Sampling: The trained decoder can sample deformations in addition to performing registration.
III. EXPERIMENTS
Experiments evaluate 3-D cardiac intra-subject registration and latent deformation representations using cine-MRI from healthy and pathological subjects.
- Experimental setting: The evaluation registers end-diastole frames to end-systole frames in cardiac cine-MRI, where the images contain large deformations.All experiments are performed in 3-D.
- Latent-space analysis: The learned deformation encoding is evaluated by visualizing the latent space and transporting encoded deformations between patients.
- Dataset: The study uses 334 ED-ES frame pairs, with 234 cases for training and 100 ACDC cases for testing.The testing set includes segmentation and disease-class information and contains healthy and pathological data.
2) Implementation details:
The study implements a multi-scale probabilistic registration system and evaluates it against established methods using intensity, anatomical-overlap, distance, and deformation-regularity measures.
- Implementation details: The network uses four encoding convolutional layers, three decoding deconvolutional layers, and two convolutional layers per scale.Each scale also includes Gaussian smoothing before exponentiation and spatial transformation.
- Implementation details: The latent code is set to d = 32 as a compromise between registration quality and generalizability, with approximately 420k trainable parameters.LeakyReLU activations and L2 weight decay of 1 ∗10−4 are used throughout most layers.
- Registration setup: The method is compared with LCC-demons, SyN, and VoxelMorph, and its three-scale objective is evaluated against a single-scale objective.VoxelMorph uses the same augmentation techniques, while SyN parameters are manually tuned on training images.
- Evaluation: Registration is assessed with RMSE, DICE, 95%-tile Hausdorff distance, and mean gradient magnitude of the determinant of the Jacobian.The Jacobian-gradient measure is used to quantify smoothness differences between deformation fields.
- Results: Our S3 achieves the best DICE performance, while both proposed variants significantly improve Hausdorff distances on the dataset.The three-scale method differs significantly from every other method across all reported metrics at p < 0.001; S1 is not significantly different from LCC-demons for DICE.
4) Deformation encoding:
The learned latent space is analyzed through dimensionality, disease clustering, deformation synthesis, and transport between pathological and healthy subjects.
- Latent vector size: Registration accuracy rises with latent dimensionality before reaching a plateau, and d = 32 is selected as a trade-off between accuracy and latent-space size.A relatively small latent size of d = 8 already achieves competitive accuracy.
- Disease distribution: 83% classification accuracy is obtained for five cardiac disease classes using eight CCA components and an SVM with 10-fold cross-validation.The 32-dimensional latent codes are projected into a two-dimensional CCA space for visualization, where test cases cluster by class.
- Generative latent space: Sampling along the two largest principal components generates deformations whose first component controls larger changes and whose second focuses on smaller changes.The experiments suggest that nearby latent codes correspond to deformations with similar characteristics.
- Deformation transport: The framework transports pathological deformations to healthy subjects by reusing predicted deformation codes without inter-subject registration.The procedure first predicts a pathological deformation and then applies its latent code with a healthy subject’s end-diastolic frame.
- Deformation transport: The transported deformation predictions produce smaller absolute ejection-fraction differences than the pole ladder for DCM cases and similarly close differences for HCM cases.The study also reports that demons underestimate HCM ejection fractions relative to the ACDC specification.
IV. DISCUSSION AND CONCLUSIONS
The framework learns a low-dimensional probabilistic deformation model for accurate registration and deformation analysis, with multi-scale estimation improving registration and deformation regularity. Its latent encoding supports deformation transport and disease clustering, while the Gaussian latent prior remains an open assumption.
- IV. DISCUSSION AND CONCLUSIONS: The method learns a low-dimensional probabilistic deformation model that supports both accurate registration and deformation analysis.The framework is unsupervised and multi-scale.
- IV. DISCUSSION AND CONCLUSIONS: The model can generate deformations from a single image and impose arbitrarily smooth diffeomorphic deformations through regularization and an exponentiation layer.The multi-scale framework provides velocities, deformation fields, and warped images at different scales.
- IV. DISCUSSION AND CONCLUSIONS: Multi-scale estimation with 32 latent dimensions significantly improved registration accuracy, while producing deformation fields smoother than the deep-learning baseline.The single-scale approach was comparable to LCC-demons and SyN; supervised masks increased DICE by only approximately 1–2%.
- IV. DISCUSSION AND CONCLUSIONS: The latent space places similar deformations near each other, enabling disease clustering and deformation transport without requiring inter-subject registration.Transport results were comparable quantitatively and qualitatively to a state-of-the-art method requiring inter-subject registration.
- IV. DISCUSSION AND CONCLUSIONS: The choice of a simple multivariate Gaussian prior for the latent space is uncertain, motivating alternatives such as mixtures of Gaussians.The authors also identify the optimal latent-vector size as an open question across applications.
LA M*
The figures compare qualitative registration outputs across cardiac cases and methods, including warped images, displacement fields, grid overlays, and Jacobian determinants. A separate comparison examines Jacobian-determinant gradients across methods and scales, while another illustrates the pole-ladder transport pipeline.
- LA M*: Figures 10–12 compare DCM, HCM, MINF, Normal, and abnormal right-ventricle cases across four registration approaches.Each case displays warped moving images, displacements, grid overlays, and Jacobian determinants.
- LA M*: The qualitative comparisons include LCC-demons, SyN, voxelmorph, and the proposed three-scale approach.The proposed method is labeled Our S3 in the figure captions.
- LA M*: Figure 13 compares the gradient of the Jacobian determinant for LCC-demons, SyN, voxelmorph, and the proposed one- and three-scale variants.The caption reports that the single-scale approach shows the most regular deformation.
- LA M*: Figure 14 presents the symbolic pipeline and a worked visualization for parallel transport using the pole ladder approach.Panel (a) shows the symbolic pipeline; panel (b) visualizes all steps for one example.