Diffusion Tensor Estimation with Uncertainty Calibration

TL;DR

This work tackles the challenge of uncertainty in diffusion tensor imaging (DTI) analysis by introducing a two-branch deep neural network that jointly estimates the diffusion tensor and its uncertainties. The model learns data-dependent (aleatoric) uncertainty via a loss term that attenuates the error for uncertain voxels and captures model uncertainty through Monte Carlo dropout, with post-hoc calibration to improve reliability. It provides a rigorous evaluation framework using calibration metrics (ENCE, AUCC) and compares against conventional CWLLS+wild bootstrap baselines, demonstrating improved accuracy with few measurements and better calibrated uncertainty under domain shift. The practical impact lies in enabling more reliable DL-based DTI analyses, including out-of-distribution detection and noise-aware interpretation of DTI biomarkers like FA and MD, ultimately aiding clinical and developmental neuroimaging studies.

Abstract

It is highly desirable to know how uncertain a model's predictions are, especially for models that are complex and hard to understand as in deep learning. Although there has been a growing interest in using deep learning methods in diffusion-weighted MRI, prior works have not addressed the issue of model uncertainty. Here, we propose a deep learning method to estimate the diffusion tensor and compute the estimation uncertainty. Data-dependent uncertainty is computed directly by the network and learned via loss attenuation. Model uncertainty is computed using Monte Carlo dropout. We also propose a new method for evaluating the quality of predicted uncertainties. We compare the new method with the standard least-squares tensor estimation and bootstrap-based uncertainty computation techniques. Our experiments show that when the number of measurements is small the deep learning method is more accurate and its uncertainty predictions are better calibrated than the standard methods. We show that the estimation uncertainties computed by the new method can highlight the model's biases, detect domain shift, and reflect the strength of noise in the measurements. Our study shows the importance and practical value of modeling prediction uncertainties in deep learning-based diffusion MRI analysis.

Figures (12)

• Figure 1: The proposed network to compute the diffusion tensor and estimation uncertainty. The residual block, denoted as RES, is shown at the bottom of the figure. ($n$ is the patch size and $m$ is the number of measurements.)
• Figure 2: Radial and axial diffusivity as a function of postmenstrual age for subjects in the dHCP (29-46 weeks) and PING datasets (9-20 years).
• Figure 3: The proposed post-hoc re-calibration method in action. This example shows the DL epistemic uncertainty for orientation of the major eigenvector on a dHCP test scan. (a) Plots of RMSE versus RMV for three training scans. The solid curve shows the isotonic regression function fitted to the data points for these three scans. This function is then used to re-calibrate the uncertainty estimations on a test scan. (b) Estimation error for the test scan. (c) Original (uncalibrated) uncertainty map for this subject and the corresponding calibration curves. (d) Uncertainty map after post-hoc re-calibration.
• Figure 4: Color FA, FA, and MD images estimated with CWLLS and DL methods on a test scan from the dHCP dataset.
• Figure 5: Estimation error and uncertainty computed with CWLLS+WBS and the propsoed DL method. $m$ is the number of measurements.