Heteroscedastic Uncertainty Estimation for Probabilistic Unsupervised Registration of Noisy Medical Images
Xiaoran Zhang* Daniel H. Pak* Shawn S. Ahn Xiaoxiao Li Chenyu You Lawrence Staib Albert J. Sinusas Alex Wong James S. Duncan
Yale University University of British Columbia
Paper | Poster | Code coming soon
Overview
|
|
We propose a heteroscedastic uncertainty estimation framework for unsupervised medical image registration. Existing methods rely on objectives (e.g. mean-squared error) that assume a uniform noise level across the image, disregarding the heteroscedastic and input-dependent characteristics of noise distribution in real-world medical images. This further introduces noisy gradients due to undesired penalization on outliers, causing unnatural deformation and performance degradation. To mitigate this, we propose an adaptive weighting scheme with a relative $\gamma$-exponentiated signal-to-noise ratio (SNR) for the displacement estimator after modeling the heteroscedastic noise using a separate variance estimator to prevent the model from being driven away by spurious gradients from error residuals, leading to more accurate displacement estimation. To illustrate the versatility and effectiveness of the proposed method, we tested our framework on two representative registration architectures across three medical image datasets. Our proposed framework consistently outperforms other baselines both quantitatively and qualitatively while also providing accurate and sensible uncertainty measures. Paired t-tests show that our improvements in registration accuracy are statistically significant. |
Motivation
|
|
Previous unsupervised image registration frameworks utilize an objective, such as mean-squared error, that assumes homoscedastic noise across an image as shown in the left figure. This does not reflect the heteroscedastic and input-dependent characteristics of noise in real-world medical imaging data. To address this issue, we propose a heteroscedastic uncertainty estimation scheme as shown in the right figure to adaptively weigh the data-fidelity term accounting for the non-uniform variations of noise across the image. |
Registration accuracy
|
|
|
Qualitative evaluation of the registration accuracy via segmentation warping for all datasets (top two rows: Voxelmorph architecture, bottom two rows: Transmorph architecture). Our method in the last column (overlayed with ground truth ES myocardium segmentation in yellow) predicts more natural and accurate deformations compared to baselines, evidenced by better matching with the ground-truth labels, smoother contour edges, and locally consistent myocardial region. |
Qualitative uncertainty evaluation
|
|
Qualitative evaluation of the registration accuracy via segmentation warping for all datasets (top two rows: Voxelmorph architecture, bottom two rows: Transmorph architecture. Our method in the last column (overlayed with ground truth ES myocardium segmentation in yellow) predicts more natural and accurate deformations compared to baselines, evidenced by better matching with the ground-truth labels, smoother contour edges, and locally consistent myocardial region. |
Quantitative uncertainty evaluation - sparsification error
|
|
|
Estimated $\hat{\sigma}_I^2$ and the corresponding weighting map of our proposed framework under Voxelmorph architecture (top row: ACDC; bottom row: CAMUS). Our predictive $\hat{\sigma}_I^2$ qualitatively makes sense since it correlates with the visual correspondence between the fixed image (first column) and the reconstructed image (second column) warped using estimated displacement. Our proposed weighting map also accurately reflects the relative importance between image intensity and predictive variance based on signal-to-noise, leading the displacement estimator to a better performance. |
Incorporating displacement uncertainty
|
|
Estimated $\hat{\sigma}_z^2$ of our proposed method framework under Voxelmorph architecture \cite{balakrishnan_voxelmorph_2019} compared with Voxelmorph-diff \cite{dalca_unsupervised_2019}. The second column shows the warped images of both methods using the displacement predicted. We overlay our predicted $\sigma_z^2$ on estimated displacement $z$ in the third column where red indicates higher uncertainty and blue indicates lower. Our estimated $\sigma_z^2$ is able to capture the randomness more accurately in $z$ together with a better registration performance than Voxelmorph-diff when comparing the second column to the fixed image due to our proposed adaptive signal-to-noise weighting strategy. |
Citation
|