Heteroscedastic Uncertainty Estimation Framework for Unsupervised Registration

TL;DR

This work addresses the problem of noise heterogeneity in unsupervised medical image registration by introducing a probabilistic framework that jointly learns a displacement estimator and a per-voxel variance estimator. It employs a collaborative training strategy with a novel gamma-exponentiated relative SNR weighting to adaptively attenuate noisy regions during optimization, improving registration accuracy across multiple architectures and datasets. Key contributions include a formal analysis of naive heteroscedastic approaches, the proposed cooperative uncertainty estimation, and validation via extensive experiments showing statistically significant gains and sensible uncertainty maps. The framework is plug-and-play, applicable to diverse cardiac imaging modalities, and offers practical uncertainty quantification with potential for data-driven objective enhancements in future work.

Abstract

Deep learning methods for unsupervised registration often rely on objectives that assume a uniform noise level across the spatial domain (e.g. mean-squared error loss), but noise distributions are often heteroscedastic and input-dependent in real-world medical images. Thus, this assumption often leads to degradation in registration performance, mainly due to the undesired influence of noise-induced outliers. To mitigate this, we propose a framework for heteroscedastic image uncertainty estimation that can adaptively reduce the influence of regions with high uncertainty during unsupervised registration. The framework consists of a collaborative training strategy for the displacement and variance estimators, and a novel image fidelity weighting scheme utilizing signal-to-noise ratios. Our approach prevents the model from being driven away by spurious gradients caused by the simplified homoscedastic assumption, leading to more accurate displacement estimation. To illustrate its versatility and effectiveness, we tested our framework on two representative registration architectures across three medical image datasets. Our method consistently outperforms baselines and produces sensible uncertainty estimates. The code is publicly available at \url{https://voldemort108x.github.io/hetero_uncertainty/}.

Figures (7)

• Figure 1: Left: We propose a heteroscedastic uncertainty estimation scheme to adaptively weight the data-fidelity term accounting for the non-uniform variations of noise across the image. Right: Overview of our proposed method. The noise variance estimator uses a U-Net backbone that takes reconstructed frame $\hat{I}_f$ along with frame $I_f$ to predict the heteroscedastic variance for the noise in \ref{['eq:sigma_I_noise']}.
• Figure 2: Qualitative evaluation of the registration accuracy via segmentation warping for all datasets (top two rows: Voxelmorph architecture balakrishnan_voxelmorph_2019, bottom two rows: Transmorph architecture chen_transmorph_2022). Our method in the last column (overlayed with ground truth (GT) ES myocardium label in yellow) predicts more natural and accurate deformations compared to baselines, evidenced by better matching with the GT, smoother contour edges, and locally consistent myocardial region.
• Figure 3: Left: Estimated $\hat{\sigma}_I^2$ and the corresponding weighting map of (top row: ACDC bernard_deep_2018; bottom row: CAMUS leclerc_deep_2019). Right: Sparsification error plots of $\hat{\sigma}_I^2$. Both plots are from our proposed framework under Voxelmorph architecture balakrishnan_voxelmorph_2019
• Figure 4: Comparison of $\hat{\sigma}_z^2$ between our vxm-based framework and vxm-diff dalca_unsupervised_2019.
• Figure 5: Qualitative evaluation for our private 3D Echo dataset on voxelmorph architecture. We extract cross-sectional slices from the 3D volume for visualization. We overlay ground truth segmentation in yellow for comparison.