Uncertainty Estimation in Radiation Dose Prediction U-Net

Abstract: The ability to quantify uncertainties associated with neural network predictions is crucial when they are relied upon in decision-making processes, especially in safety-critical applications like radiation therapy. In this paper, a single-model estimator of both epistemic and aleatoric uncertainties in a regression 3D U-net used for radiation dose prediction is presented. To capture epistemic uncertainty, Monte Carlo Dropout is employed, leveraging dropout during test-time inference to obtain a distribution of predictions. The variability among these predictions is used to estimate the model’s epistemic uncertainty. For quantifying aleatoric uncertainty quantile regression, which models conditional quantiles of the output distribution, is used. The method enables the estimation of prediction intervals of a user-specified significance level, where the difference between the upper and lower bound of the interval quantifies the aleatoric uncertainty. The proposed approach is evaluated on two datasets of prostate and breast cancer patient geometries and corresponding radiation doses. Results demonstrate that the quantile regression method provides well-calibrated prediction intervals, allowing for reliable aleatoric uncertainty estimation. Furthermore, the epistemic uncertainty obtained through Monte Carlo Dropout proves effective in identifying out-of-distribution examples, highlighting its usefulness for detecting anomalous cases where the model makes uncertain predictions.