Hierarchical uncertainty estimation for learning-based registration in neuroimaging

2024-10-11

arXiv.org

AI Summary

Weak uncertainty modeling in neuroimaging registration—particularly the neglect of spatial structural characteristics by existing methods (e.g., Monte Carlo Dropout)—limits reliability in clinical applications. To address this, we propose the first hierarchical uncertainty propagation framework tailored for MRI registration. Our method models voxel-wise local uncertainty via Gaussian distributions, with theoretical justification, and introduces a hierarchical probabilistic propagation mechanism that systematically transfers both epistemic and aleatoric uncertainties—from local spatial cognition through global deformation fields to downstream tasks (e.g., segmentation, statistical mapping). The framework integrates Monte Carlo Dropout for ablation analysis, posterior transformation sampling, and a deep registration network. Experiments demonstrate: (i) significantly stronger correlation between estimated uncertainty and ground-truth registration error versus baselines; (ii) improved registration accuracy in brain MRI when leveraging uncertainty-aware optimization; and (iii) robust uncertainty propagation to downstream tasks, enhancing clinical decision reliability.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Relational Probabilistic Models
• Multiagent Systems: Multiagent Systems under Uncertainty

Application Category

• Responsible Web: Machine-in-the-loop, human agency and autonomy
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Over recent years, deep learning based image registration has achieved impressive accuracy in many domains, including medical imaging and, specifically, human neuroimaging with magnetic resonance imaging (MRI). However, the uncertainty estimation associated with these methods has been largely limited to the application of generic techniques (e.g., Monte Carlo dropout) that do not exploit the peculiarities of the problem domain, particularly spatial modeling. Here, we propose a principled way to propagate uncertainties (epistemic or aleatoric) estimated at the level of spatial location by these methods, to the level of global transformation models, and further to downstream tasks. Specifically, we justify the choice of a Gaussian distribution for the local uncertainty modeling, and then propose a framework where uncertainties spread across hierarchical levels, depending on the choice of transformation model. Experiments on publicly available data sets show that Monte Carlo dropout correlates very poorly with the reference registration error, whereas our uncertainty estimates correlate much better. Crucially, the results also show that uncertainty-aware fitting of transformations improves the registration accuracy of brain MRI scans. Finally, we illustrate how sampling from the posterior distribution of the transformations can be used to propagate uncertainties to downstream neuroimaging tasks. Code is available at: https://github.com/HuXiaoling/Regre4Regis.

Problem

• Estimating uncertainty in deep learning-based image registration.
• Propagating uncertainties from local to global transformation models.
• Improving registration accuracy in brain MRI scans using uncertainty-aware fitting.

Innovation

• Hierarchical uncertainty propagation in image registration
• Gaussian distribution for local uncertainty modeling
• Uncertainty-aware transformation fitting improves MRI accuracy