Dual Heteroscedastic Uncertainty Estimation for Probabilistic Unsupervised Volumetric Registration of Noisy Medical Images

Medical images are often subject to spatially non-uniform image noise due to various factors including the imaging technique, the underlying tissue properties, and the imaging conditions. Despite this intrinsic heterogeneity, previous learning-based unsupervised image registration techniques have primarily operated under the simplified homoscedastic assumption, such as an additive Gaussian noise with constant variance across the image space. This leads to an equally weighted image fidelity loss term, which has the potential to overemphasize image noise and introduce unnatural deformations. To mitigate this, we propose a novel probabilistic unsupervised registration framework that explicitly estimates and leverages \textit{heteroscedastic} image noise in the learning process. We present a collaborative learning strategy, where we jointly train a motion estimator and a variance estimator using separate objectives that include an improved signal-to-noise ratio (SNR)-based weighting strategy. We tested our method across diverse cardiac imaging datasets, including public 2D MRI, public 2D ultrasound, and a private 3D echocardiography dataset. Our method shows improved registration performance.