Hierarchical Uncertainty Estimation for Learning-based Registration in Neuroimaging

TL;DR

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 is proposed.

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.

Figures (6)

• Figure 1: Overview of the training strategy of the proposed method.
• Figure 2: Coordinate prediction error in millimeters ($mm$) (Error) and estimated variance in $mm^2$ (Var) for the epistemic (epis.) uncertainty (a, b) and for the aleatoric (alea.) uncertainty (c, d).
• Figure 3: The top row shows samples from the B-spline transformation with their variance and sample-to-sample differences displayed on the second row (from left to right). Similarly, the third row shows samples from the Demons transformation with their variance and sample-to-sample differences shown on the last row (from left to right). Samples and their differences. Diff: 2$\rightarrow$1 means the difference between sample 2 and sample 1, and the same applies to others.
• Figure 4: The heat map of uncertainty from B-spline and Demons. Note the B-spline coefficients are upsampled to the image size for visualization.
• Figure 5: Example samples from the B-spline transformation (top row, (b-e)) and the Demons transformation (bottom row (h-k)), along with the input scan (a) and the ground-truth segmentation (g). The last column (f, l) show the entropy of the sampled segmentation.