Quality control in radiotherapy-treatment planning using multi-task learning and uncertainty estimation

Felix J.S. Bragman, Ryutaro Tanno, Zach Eaton-Rosen, Wenqi Li, David J. Hawkes, Sebastien Ourselin, Daniel C. Alexander, Jamie R. McClelland, M. Jorge Cardoso

Abstract: Multi-task learning is ideally suited for MR-only radiotherapy planning as it can jointly simulate a synthetic CT (synCT) scan - a regression task - and an automated contour of organs-at-risk - a segmentation task - from MRI data. We propose to use a probabilistic deep-learning model to estimate respectively the intrinsic and parameter uncertainty. Intrinsic uncertainty is estimated through a heteroscedastic noise model whilst parameter uncertainty is modelled using approximate Bayesian inference. This provides a mechanism for data-driven adaptation of task losses on a voxel-wise basis and importantly, a measure of uncertainty over the prediction of both tasks. We achieve state-of-the-art performance in the regression and segmentation of prostate cancer scans. We show that automated estimates of uncertainty correlate strongly in areas prone to errors across both tasks, which can be used as mechanism for quality control in radiotherapy treatment planning.

Keywords: Deep learning, multi-task learning, probabilistic modelling, image synthesis, semantic segmentation

Author Affiliation: University College London, King's College London, National University of Singapore