Refining Bayesian neural networks via Stein variational gradient descent for uncertainty-aware cardiac registration

Bayesian neural networks (BNNs) enable learning a probability distribution over network weights and quantifying uncertainty for medical image registration. In this paper, we propose a Bayesian variational autoencoder (BVAE)-based image registration framework, termed BVAE, that simultaneously provides both epistemic and aleatoric uncertainty with improved accuracy by jointly estimating the posterior distributions of network weights and latent variables of deformation vector fields via variational inference. Moreover, the prior distribution over network weights is learned as an optimal moment-matching prior, further improving posterior expressiveness. A signal-to-noise ratio is defined based on epistemic uncertainty and leveraged to reduce it. Weight particles sampled from the BNN are projected into a lower-dimensional subspace using principal component analysis (PCA) and refined by Stein Variational Gradient Descent (SVGD). Both uncertainty calibration and predictive performance are further improved through SVGD-based refinement. Extensive experiments conducted on public cardiac image registration benchmarks demonstrate that our BVAE outperforms state-of-the-art registration models, particularly in scenarios with limited training and cross-domain data, while providing aleatoric uncertainty and well-calibrated epistemic estimates for cardiac registration applications. The source code is publicly available at: https://github.com/nanase0629/SVGD-BNN-Cardiac-Registration.