Heteroscedastic Uncertainty Estimation Framework for Unsupervised Registration

Deformable image registration solves for a dense pixel-wise displacement map that aligns one image to another. It is a key medical image analysis technique that enables disease progression monitoring and surgical guidance [8, 18, 24, 31]. Classical methods (e.g. elastic-type models [13], free-form deformation with b-splines [21], diffeomorphic models [1]) are often computationally expensive and impractical for large-scale analyses. Thus, numerous works have proposed training a neural network to accelerate the test-time prediction of the displacement map [3, 5, 9, 11, 32, 30]. Such frameworks commonly utilize an unsupervised objective by minimizing the mean-square error (MSE) between the warped and target images. By using this objective, the existing frameworks assume an additive homoscedastic Gaussian image noise, i.e. $\epsilon\sim\mathcal{N}(0,\sigma^{2})$, with a constant noise variance across the spatial domain. However, this assumption is problematic for medical images, where the noise is intrinsically heteroscedastic and input-dependent (e.g. MRI [7, 26] or ultrasound [19, 29]). The non-uniform noise variance across the image space (Fig. 1, left) is due to multiple factors including changes in anatomical structures or patient motion artifacts. The simplified homoscedastic assumption disregards such variations in noise levels and results in undesired penalization of false noise-induced outliers, causing unnatural and inaccurate deformations.

To mitigate this, we propose a probabilistic heteroscedastic noise modeling framework for unsupervised image registration (Fig. 1, right). Our method involves two modules: (1) a standard displacement estimator and (2) a variance estimator that predicts the input-specific heteroscedastic noise for each image pair. For proper utilization of the predicted noise, we introduce a novel adaptive weighting strategy based on the relative $\gamma$-exponentiated signal-to-noise ratio (SNR). This enables the successful convergence of our collaborative training strategy, where the two estimators improve upon each other via information exchange and loss calibration. We validated the effectiveness and versatility of our proposed framework with extensive experiments using two representative neural network architectures and three cardiac datasets. Our proposed framework can be operated in a plug-and-play manner, and provides sensible heteroscedastic uncertainty measures to reflect spatially varying noise. This work may help provide a new perspective in data-driven input noise modeling for deep learning-based unsupervised registration.

Our contributions include (1) We first analyze a naive approach of applying heteroscedastic noise modeling to unsupervised registration and identify the pitfalls of such an approach. (2) From which, we propose a probabilistic framework for data-driven estimation of heteroscedastic uncertainty for unsupervised registration. (3) We introduce an adaptive $\gamma$-exponentiated relative SNR weighting strategy for proper loss calibration. (4) We demonstrate the effectiveness and versatility of our method on two representative registration architectures across three datasets, demonstrating consistent statistically significant improvements over the baselines while providing sensible uncertainty maps.

Let $I:\Omega\rightarrow\mathbb{R}$ be an image and $\Omega\subseteq\mathbb{R}^{d}$ its spatial domain. Unsupervised image registration aims to find the alignment between a moving image $I_{m}$ and a fixed image $I_{f}$, without ground truth information about warping displacement field $z:\Omega\rightarrow\mathbb{R}^{d}$. For each image pair $(I_{m},I_{f})$, we assume $I_{m}$ is a fixed parameter and $I_{f}$ is a noisy observation of the warped image $\hat{I}_{f}=I_{m}\circ z$. Then, $z$ is a sample from the posterior distribution $p(z|I_{f};I_{m})$.

With this formulation, the main task is finding the displacement posterior $p(z|I_{f};I_{m})$ that will maximize the data likelihood $p(I_{f};I_{m})$. Since directly solving for $p(z|I_{f};I_{m})$ is intractable, we apply the variational approach to approximate it using a neural network $q_{\theta}(z|I_{f};I_{m})$. To minimize the discrepancy between approximate and real posterior distributions, we compute the KL-divergence and arrive at the negative evidence lower bound:

For the warped image distribution $p(I_{f}|z;I_{m})$, we assume an additive zero-mean heteroscedastic Gaussian noise with spatially varying variance $\sigma_{I}^{2}:\Omega\rightarrow\mathbb{R}^{+}$ to incorporate input-dependent heteroscedasticity

Similar to [3, 12], we assume pixel-wise independence, isotropic Gaussian for the displacement posterior $q_{\theta}(z|I_{f};I_{m})=\mathcal{N}(z;\mu_{z},\sigma_{z}^{2}\mathbb{I})$, and uniform Gaussian for the displacement prior $p(z)=\mathcal{N}(0,\mathbb{I})$. Then, we arrive at the preliminary loss to be minimized for each image pair

where $\hat{I}_{f}=I_{m}\circ z$. $\sigma_{I}^{2}$ is estimated by our variance estimator as $\hat{\sigma}_{I}^{2}=h_{\phi}(I_{f},\hat{I}_{f})$, using the fixed and moved image pair as input. $z$ is estimated by our displacement estimator as $\hat{z}=q_{\theta}(I_{m},I_{f})$, using the moving and fixed pair as input.

We propose a heteroscedastic uncertainty estimation framework for probabilistic unsupervised image registration to adaptively weight the displacement estimation with relative $\gamma$-exponentiated signal-to-noise, which improves registration performance from the commonly used homoscedastic assumption while also providing accurate and sensible uncertainty measures. Our proposed framework consists of a displacement and a variance estimator, optimized under an alternating collaborative strategy. We demonstrate the effectiveness and versatility of our proposed framework on two representative registration architectures across diverse cardiac datasets and show consistent statistically significant improvements over other baselines. Though our proposed framework is promising, it still relies on a manually crafted adaptive map on the data fidelity term, which might not be able to fully reflect the complicated characteristics of large-scale real-world data. Future work will aim to explore a more data-driven objective with further validation of clinical datasets for more potential impact.