Variational image registration with learned prior using multi-stage VAEs.

Area of Science:

• Medical Imaging
• Machine Learning
• Computer Vision

Background:

• Variational Autoencoders (VAEs) are used in medical image registration due to their inherent uncertainty quantification.
• A key limitation of VAEs is the inadequate regularization from simple priors, causing a mismatch with the variational posterior.

Purpose of the Study:

• To propose a multi-stage VAE approach for learning an optimal prior, specifically the aggregated posterior, to enhance medical image registration.
• To address the mismatch between variational posterior and prior in VAEs for improved regularization.

Main Methods:

• A multi-stage VAE framework is developed to learn the aggregated posterior as an optimal prior.
• A factorized telescoping classifier estimates density ratios for accurate KL divergence calculation.
• A low-rank covariance matrix is learned to capture correlations between latent variables, reducing deformation field uncertainty.

Main Results:

• The proposed method effectively estimates high-dimensional aggregated posteriors common in medical image registration.
• Analysis revealed an optimal level of factorization for KL divergence and registration accuracy.
• Learned covariance matrices improved the handling of latent variable relationships and reduced deformation uncertainty.

Conclusions:

• The multi-stage VAE with an learned optimal prior significantly enhances medical image registration accuracy.
• The method demonstrates superior performance in negative log-likelihood (NLL) and registration accuracy across multiple datasets compared to existing approaches.