Deep Evidential Learning for Radiotherapy Dose Prediction

In the standard supervised learning formalism, typically the model is trained via minimizing a loss function such as the mean squared error so that the neural networks’ weights converge to at least a local minimum point, yielding a good accuracy defined in terms of the error term implied by the loss function. This standard approach does not directly give any estimates of model or data noise uncertainties. In a more refined approach, one can attempt to provide an estimate of uncertainty by the following assertion:

• •

  The target outputs can be modeled as being drawn from a probability distribution. For example, in our context of dose prediction, we let the predicted dose $y_{k}$ in voxel $V_{k}$ be described by a Gaussian probability distribution function (PDF) with mean and variance parameters $(\mu_{k},\sigma_{k})$.

• •

  We then attempt to train the model to infer the set of $(\mu_{k},\sigma_{k})$ (for every voxel $k$) by the principle of maximum likelihood estimation. Typically, this is performed by minimizing the negative log-likelihood function $\mathcal{L}$ defined as

  $$f(y_{k}|\mu_{k},\sigma_{k})=\frac{1}{\sqrt{2\pi\sigma^{2}_{k}}}e^{-\frac{(y_{k}-\mu_{k})^{2}}{2\sigma^{2}_{k}}},\,\,\,\,\,\,\mathcal{L}(\vec{\mu},\vec{\sigma})=-\log\left(\Pi_{k}f(y_{k}|\mu_{k},\sigma_{k})\right).$$

  (1)

  For an ensemble of models being trained upon the same dataset to yield estimates of the parameters $(\mu_{k},\sigma_{k})$, one would obtain a spectrum of differing values of $(\mu_{k},\sigma_{k})$, the distribution being stochastic in nature due to distinct initial weight distributions, and their sensitivity to the random noise present in the training dataset.

The Deep Evidential Learning framework of [30] and [29] posits that there is another PDF which describes the distribution of the Gaussian mean and variance parameters. This can be regarded as a ‘higher-order’ PDF that in turn describes the distribution of $(\mu_{k},\sigma_{k})$ of the ‘first-order’ PDF $f(y|\mu,\sigma)$ in eqn. (1). In Bayesian theory, there is a natural mathematical entity that enacts this role – the prior probability distribution that treats $(\mu_{k},\sigma_{k})$ as random rather than deterministic variables. In [29], various model training and simulations were performed by taking the prior distribution to be a product of a Gaussian distribution (for $\mu_{k}$) and an inverse-gamma distribution (for $\sigma_{k}$). Each of them is separately a two-parameter distribution. Suppressing the voxel index $k$, we have

where $\Gamma^{-1}(x)$ is the inverse gamma distribution. The overall prior distribution is a four-parameter normal-inverse-gamma distribution defined as the product of the two PDFs above.

Given this prior distribution, we can take expectation values and other statistical moments with respect to it to compute mean and variances. In the simplest form, Deep Evidential Learning implies attaching a final feedforward layer (e.g. of the convolutional type) to an existing model architecture such that we have $\{\alpha,\beta,\nu,\gamma\}$ as the eventual model outputs.

The mean prediction, epistemic and aleatoric uncertainties are then defined with respect to $\mathcal{P}$ as follows.

where $f$ is the ‘first-order’ PDF of eqn. (1). It is in this sense that $\mathcal{P}$ is the ‘higher-order’ PDF providing the description of $\mu,\sigma$ of $f$ as random variables. We note that the aleatoric and epistemic uncertainties are obtained after integrating over $\mu,\sigma$ using eqn. (4). The factorized form of $\mathcal{P}$ enables an analytic expressions for each of these uncertainties. In Deep Evidential Learning, the model output parameters are $\{\alpha,\beta,\nu,\gamma\}$. From these parameters, we can then obtain the dose prediction (given by $\gamma$) as well as the uncertainty estimates $U_{a},U_{e}$ for each voxel.

We now apply such an uncertainty quantification framework in the context of a dose prediction model which estimates the dose delivered to each CT voxel based on an input set of CT images and various binary-valued masks representing the radiotherapy’s target regions and the organs-at-risk. The OpenKBP dataset of [7, 14] was used as our source, and a vanilla 3D U-Net was employed as the backbone architecture.

In [29], the first-order PDF pertaining to the model output $y$ was assumed to be the normal distribution $\mathcal{N}(y|\mu,\sigma)$, with a 4-parameter normal-inverse-gamma distribution ( eqn. (4) ) adopted to be the higher-order evidential distribution. For our case, a crucial caveat is that the radiation dose value is definitively constrained within a finite interval, and not unbounded. To incorporate this physical nature of the radiation dose, we first define a dimensionless dose parameter $y=y(D_{p})=\frac{0.9D_{p}+10}{100}$ where $D_{p}$ is the physical dose in units of Gy. We then pass $y$ to its logit representation $\log\left(\frac{y}{1-y}\right)$ which we take as the form of output for the neural network (rather than the physical dose $D_{p}$). Ablation experiments using the validation dataset were used to determine the coefficients in the linear map $y=y(D_{p})$ such that the logit representation yielded the optimal model performance. Our choice of the linear map $y=y(D_{p})$ also avoided the singular points $\{0,1\}$ of the logit function by construction.

Formally, our choice of the output data representation implies that the first-order PDF for each voxel is related to the logit-normal distribution more precisely instead of a Gaussian. Its full form reads

This is a close cousin of the normal distribution, and has a bounded domain $y\in(0,1)$ as the support. The parameters $\mu,\sigma$ are now dimensionless parameters defined as the Gaussian mean and standard deviation of the logit of the (dimensionless) dose $y$. For their higher-order evidential distributions, we can still follow the procedure of [29], i.e. we adopt a normal distribution for $\mu$ and an inverse-gamma distribution for the variance to yield a normal-inverse-gamma distribution for (6). Also, similar to the approach in [29], we can derive the aleatoric and epistemic uncertainties. In their original forms as defined in eqn. (5), they are dimensionless variances associated with the logit representation. To convert them to the corresponding dimensionful ones in units of Gy², we first invoke the standard approximation technique of relating the uncertainty in the $L(y)$ (logit of $y$) to that in $y$ itself, before using the linear map $y=y(D_{p})$ to deduce the physical aleatoric and epistemic uncertainties in units of Gy². For clarity, we refer to them separately as $U_{alea},U_{epis}$ which are related to $U_{a},U_{e}$ in eqn. (5) via $U_{alea,epis}=\left(\frac{100}{0.9}y(1-y)\right)^{2}U_{a,e}$.

Following [29], the primary loss function can be taken to be the maximum likelihood loss defined as the marginal likelihood obtained when we integrate over the first-order mean $\mu$, and variance $\sigma^{2}$. After some algebra, we find that we obtain the likelihood function

where the expression lying on the right of the term $1/(y(1-y))$ can be identified as the Student’s t distribution $\text{St}\left(\log(\frac{y}{1-y});\gamma,\frac{\beta(1+\nu)}{\nu\alpha},2\alpha\right)$. Unfortunately, we found that in this form, the loss function did not lead to stable learning curves that could generate uncertainty heatmaps that correlated with prediction errors. After some experimentation, we found that making the following refinements to eqn. (8) significantly enhanced the effectiveness of the model:

• •

  removal of the factor $1/y(1-y)$ from eqn. (8),

• •

  passing the negative log-likelihood function through a regularizing positive-definite activation function (our choice: standard logistic function),

• •

  adding a mean-squared-error loss function term.

These refinements led to the eventual form of our loss function being as follows.

where $\lambda_{KL},\lambda_{mse}$ are crucial hyperparameters of which optimal values for our purpose will be discussed in Sec. 3.2. The first regularization term $\mathcal{L}_{reg}=\left|\log\left(\frac{y}{1-y}\right)-\gamma\right|(2\nu+\alpha)$ preceded by $\lambda_{KL}$ was proposed in [29] to penalize statistical evidence for incorrectly labeled terms. We note in passing that in [29], it was argued that the prior (the function $\mathcal{P}(\mu,\sigma|\alpha,\beta,\nu,\gamma)$ in eqn. (4) ) has a singular limit when we take parameter values corresponding to ‘zero evidence’, and that softening this limit by introducing finite yet small values for $\alpha=1+\epsilon$, $\nu=\epsilon$ did not appear to be effective. We found this to be the case as well. Another possibility that we explored is using Jeffrey’s prior (which in this case involves an improper uniform prior for the mean of the model output). But this turned out to be similarly ineffective relative to the choice adopted in [29].

Among the many papers devoted to the subject of dose prediction models (see e.g. [5] for an extensive review and compilation), to our knowledge, there has been only two which directly delved into the notion of uncertainty quantification frameworks: [8] and [28]. These papers contained highly interesting results which we would like to briefly discuss in relation to our work.

In [8], the authors studied the frameworks of MC Dropout and bootstrap aggregation which is essentially a Deep Ensemble method where each model is trained on only a portion of the entire training dataset. Like in this work, they used the OpenKBP Challenge dataset of [14] for validating the performance of each uncertainty quantification framework, concluding that Deep Ensemble yielded a lower mean absolute error (MAE) while showing better correlation between uncertainty and prediction error. In particular, the authors proposed that with the aid of an additional scaling factor, the models generated uncertainty heatmaps and DVH confidence intervals which appeared reasonable. This scaling factor was defined differently for each region of interest (ROI) and was argued to be necessary to bring the uncertainty values into a more ‘interpretable scale’ as the raw ones were only providing ‘relative measures’. For each set of points belonging to a ROI in the validation dataset, this ROI-dependent factor was defined simply as the standard deviation of the ratio between the prediction error and the uncertainty value (see eqn. (3),(4) of [8]). In our work, we did not incorporate such an empirical scaling factor in our definitions of various uncertainties, including our final uncertainty heatmaps and the DVH confidence intervals. For Deep Evidential Learning where the predictive variance can be expressed precisely in terms of learned output variables of the model, in principle, it is not clear to us how an additional empirical scaling factor should be justified. By construction, this rescaling factor naturally enhances the correlation between uncertainty and error since the uncertainty values themselves are part of the factor’s definition, assuming that it generalizes well for unknown (e.g. testing) data. Conversely, we took the resulting average value of such a ratio (across all voxels) to assess whether the uncertainty distribution generated by the model was reasonable in terms of the overall scale of magnitude. Through the numerical values (see Table 1) and the DVH certainty intervals in Fig. 8, we found that this was indeed the case.

The work in [28] proposes the use of Gaussian mixture models of which parameters are the neural network’s outputs, and with U-Net as the backbone architecture. Such a mixture density network is very similar in form to our Deep Evidential Learning model since the output variables of our model are also parameters of a PDF. In our case, the PDF is a higher-order normal-inverse-gamma distribution whereas in [28], the parameters are the means and variances of the component normal distributions (see eqn. (1)) and the coefficients defining their linear combination. Each component of the Gaussian mixture pertained to a treatment protocol representing specific priorities (e.g. OAR sparing over target coverage, etc.), and the relative values of the variances quantified the trade-offs between these protocols. The cross-entropy loss function associated with the cumulative distribution function (CDF) of the predicted dose was then taken as part of the objective function for generating a dose mimicking model. In [28], deliverable treatment plans were the endpoints of their dose prediction algorithm, with the optimization problem solved by RayStation’s native sequential quadratic programming solver and final dose computed via a collapsed cone algorithm [28]. Although we noted that there was no scrutiny of the correlation between uncertainty and error, the authors of [28] (and its follow-up work [27]) notably pointed out that probabilistic dose prediction models such as [8] and ours could play a crucial role in their pipeline since our model would yield dose CDFs that can be used to define the cross-entropy loss function as part of the dose-mimicking model’s objective function. It would thus be interesting to extend results of our work here by integrating our probabilistic dose predictions into a pipeline like [28] where deliverable treatment plans with specified machine parameter settings can be anticipated as final outcomes.

The OpenKBP Challenge dataset of [7, 14] is devoted towards establishing an open framework for the development of plan optimization models for knowledge-based planning (KBP) in radiotherapy. It is an augmented variation of real clinical data constructed as follows: medical images were taken from several institutions available on The Cancer Imaging Archive (TCIA) [37] which hosts open-source data that has undergone de-identification compliant with DICOM standard.^ivivivIn particular, all metadata and Protected Health Information (PHI) has been removed from the datasets provided in OpenKBP Challenge [7, 14]. After removing highly incomplete imaging datasets to obtain a final competition dataset comprising of 340 patients, synthetic radiation plans were generated for each of them using a published automated planning method [38]. This final dataset is then split randomly by the organizers of OpenKBP Challenge to be as follows: 200 (training), 40 (validation), 100 (testing). In our work, we followed this decomposition of data. As in a conventional machine learning pipeline, the validation dataset was used for tuning of hyperparameters whereas the testing dataset was used for reporting various results in our work. We refer the reader to [39, 40] for more pedagogical expositions of data-splitting practices and algorithms.

The radiotherapy plans were delivered using nine equidistant coplanar beams at various angles with a 6 MV step-and-shoot intensity-modulated-radiation-therapy (IMRT) in 35 fractions. The organizational team of the Challenge used the IMRTP library from A Computational Environment for Radiotherapy Research implemented in MATLAB [41] to generate the dose deposited at each voxel. These dose distributions were from fluence-based treatment plans with similar degrees of complexity [14, 42].

We used a $(128,128,128,11)$-dimensional input representation consisting of (i)CT in Hounsfield units(clipped to be within [0, 4095] before being normalized to [0,1]), (ii)structure masks of three planning target volumes (PTV) and seven organs-at-risk (OAR) regions : each is a Boolean tensor labeling any voxel contained within the respective structure. The OARs are brainstem, spinal cord, right and left parotids, larynx, mandible and esophagus. The three PTV regions are: PTV56, PTV63, PTV70 which are targets that should receive 56 Gy, 63 Gy and 70 Gy of radiation dose respectively.

Since our primary goal is to examine the uncertainty estimation aspects of the model, in this work, we adopt a simple vanilla 3D U-Net as the backbone architecture upon which we insert two additional layers so that they are compatible with the framework of Deep Evidential Learning. U-Nets and their variants have featured heavily in many deep-learning-guided tasks related to radiotherapy dose prediction [8, 14]. They are essentially convolutional neural networks equipped with an encoder-decoder architecture that allows effective learning of both image features at various levels of resolution (see e.g. [43, 44] for reviews on U-Net and [45, 46] for pedagogical introductions to convolutional neural networks).

Our U-Net has the following structure (see Fig. 1):

• •

  it takes in a 11-channel input of size $128\times 128\times 128$ voxels,

• •

  each downsampling level consists of two convolutional layers each with a $(3\times 3\times 3)$ kernel, ReLU activation and equipped with a dropout unit, followed by maximum pooling with kernel size $(2\times 2\times 2)$,

• •

  the downsampling operation is performed 4 times, with the dropout rates increasing consecutively as $\{0.10,0.15,0.20,0.25\}$, and the number of convolutional filters for each layer being $\{16,32,64,128\}$ respectively,

• •

  the bottleneck layer has 256 filters and dropout rate of 0.30,

• •

  at each of the four upsampling levels, features from the contraction path are concatenated with the corresponding upsampled features, and the dropout rates decrease similarly in the reverse order.

• •

  a 8-channel pointwise convolution is then applied followed by a final 4-channel pointwise convolution that yields the four parameters $\{\alpha,\beta,\nu,\gamma\}$ of the Bayesian prior distribution $\mathcal{P}(\mu,\sigma|\alpha,\beta,\nu,\gamma)$ of eqn. (4).

This model carries about $6\times 10^{6}$ free weight parameters. Various hyperparameters such as the dropout rate at each level, etc. were eventually adopted after running ablation experiments to determine their optimal values.

Imposing some choice of numerical bounds to $\{\alpha,\beta,\nu,\gamma\}$ in the final layer is a crucial aspect of the hyperparameter tuning process unique to this framework. The parameter $\beta$ sets the overall (dimensionless) scale to both uncertainties of which ratio is set by $\nu=U_{alea}/U_{epis}$ which has to be positive. For this dataset, we found good uncertainty-error correlation after imposing $\beta\geq 10^{-3}$. In [29], the authors imposed the constraint $\alpha>1$ which we also followed here, as for this range of values, the aleatoric and epistemic uncertainties can be conveniently described by simple analytic expressions. More broadly speaking, $\alpha>0$ is the larger admissible range for $\alpha$ as a parameter of the inverse-gamma distribution. We also found it useful to map the zero dose to a small positive value $\epsilon=0.1$ in the (dimensionless) normalized dose $D_{N}$ defined to be $D_{N}=\frac{0.9D_{p}+10}{100}$ where $D_{p}$ is the physical dose in Gy. In the loss function, we work in the logit-representation of $D_{N}$, with the predicted dose $\gamma=\log\left(\frac{D_{N}}{1-D_{N}}\right)$. Since this logistic function is monotonic, $\gamma$ is bounded from below as $\gamma\geq\log\left(\frac{\epsilon}{1-\epsilon}\right)=-\log 9$ with our choice of $\epsilon=0.1$.

For the loss function originally proposed in [29], we found that unfortunately, it led to training curves with frequent oscillations and poor eventual outcomes. As explained in detail in Sec. 2.2, we reformulated the loss function for a logistic representation of the dose variable such that one could obtain a stable implementation. Our refined loss function in eqn. (10) is characterized by the hyperparameters $\lambda_{mse},\lambda_{KL}$ which describe couplings of the loss function to a mean-squared-error (MSE) term and a KL-divergence-like regularization term respectively. Results from ablation experiments with validation dataset indicated that optimal values of these hyperparameters lie within the intervals $\lambda_{mse}\in(0.01,0.1)$, $\lambda_{KL}\lesssim 0.01$. For the results reported here for the rest of our paper, we took $\lambda_{mse}=0.05,\,\lambda_{KL}=0.01$. In Fig. 2, we plot the learning curves pertaining to the original and our refined loss functions for comparison.

Towards comparing the Deep Evidential model against more common approaches of uncertainty quantification in literature, we implemented a Deep Ensemble model and Monte-Carlo Dropout model for the same dataset. Each of the 5 neural networks used in the ensemble model shared the same 3D U-Net architecture with all hyperparameters preserved. Random weight initialization was all performed via the He-uniform initializer [47]. Similarly, for the Monte-Carlo Dropout model, the same base model was used for passing 30 forward passes with the dropout layers activated for model prediction. All models were trained for 200 epochs with a learning rate of $10^{-4}$ using Adam optimizer. Our focus in this work lies in studying aspects of uncertainty estimation, and thus our choice of a relatively simple backbone architecture equipped with only $\sim 6\times 10^{6}$ weight parameters. Compared to other more complicated setups, each of these models converged quickly in about a couple of hours while attaining a mean-absolute-error of dose prediction that laid within the top 20 of the score chart of the OpenKBP Challenge [14].

The results in Table 1 indicated that Deep Evidential Learning model yielded spectra of uncertainty values which were more correlated with prediction error, relative to MC Dropout and Deep Ensemble methods. To further scrutinize the difference among the various frameworks, we examined how error distributions varied with changing threshold values of uncertainty. Since the error distributions were generically found to be quite skewed, we used the median of the error distribution as its characterizing parameter. Fig. 3 and 4 below reveal how median error changed with increasing levels of uncertainty thresholds for the different models. Assuming that uncertainty values correlate positively with prediction error implies that one expects a monotonically increasing curve in such a plot, a trend that was indeed manifest for these models as shown in Fig. 3 and 4.

In particular, from Fig. 4, the visibly more evident linearity of the curve for Deep Evidential Learning indicated a sensitivity of epistemic uncertainty to prediction errors that was relatively more uniformly calibrated compared to Monte-Carlo Dropout and Deep Ensemble methods.

A prominent feature of Deep Evidential Learning is a clear principled approach towards distinguishing between aleatoric and epistemic uncertainties. They can be defined precisely using the higher-order PDF $\mathcal{P}$ as expressed in eqn. (5). Traditionally, aleatoric uncertainty is commonly interpreted as arising from inherent random noise in the data, in contrast to epistemic uncertainty which presumably captures uncertainty originating from model suitability and data sufficiency (see e.g. [32, 36]).

We would like to probe the level of sensitivity of various uncertainty measures to noise perturbations of the input data. In particular, we would like to identify if there would be any difference in noise-sensitivity between $U_{alea}$ and $U_{epis}$. Our choice of perturbation is the addition of Gaussian noise [48] of mean zero and standard deviation 0.5 to each CT voxel (of which intensity was preprocessed to lie within $[0,1]$). To enable a more holistic description of response to noise, we plot the empirical cumulative distribution function (eCDF) of various uncertainty measures preceding and following the noise addition.

In Fig. 5, the plot of various eCDFs collectively illustrated the differences among the uncertainty measures in their responses towards the Gaussian perturbation. Since different measures yielded distinct overall scale of uncertainty magnitudes, we normalized each with respect to the maximum value for each uncertainty type. In the limit of the Gaussian noise completely dominating over any pre-existing noisiness, and assuming that the uncertainty value of each voxel is correlated only with the noise PDF (identical for every voxel), we would expect the uncertainty eCDF to approach a symmetric distribution with mean (normalized) uncertainty at 0.5 in this limiting scenario. As a simple reference distribution, in Fig. 5, we included the line corresponding to the uniform distribution. From just visual inspection, aleatoric uncertainty distribution and that of Deep Ensemble changed more significantly relative to others. The noise-induced fractional changes in the relative entropy (measuring the Kullback-Leibler divergence between each eCDF and the uniform distribution) turned out to be $\sim 0.1$ for aleatoric uncertainty and Deep Ensemble, about ten times larger than those for epistemic uncertainty and Monte-Carlo dropout.

The various curves appeared to suggest the following.

• •

  The relative responsiveness of aleatoric (stronger) and epistemic (weaker) uncertainties are indeed compatible with their conventional interpretations as reflecting inherent data noise and model-related uncertainties respectively.

• •

  Monte-Carlo Dropout yielded an uncertainty distribution that, like epistemic uncertainty, was not sensitive to the addition of the Gaussian noise. This favored its interpretation as an epistemic uncertainty in nature.

• •

  Deep Ensemble method yielded an uncertainty distribution that was responsive to the addition of noise. Like the aleatoric uncertainty in Deep Evidential Learning, the eCDF shifted to being less long-tailed and leaning towards being more uniform.

For the Deep Ensemble method, it was explained in [34, 49] that one can often use a Gaussian mixture model as an effective description of the ensemble prediction. By Eve’s law of total variance [49], the predictive variance admits a natural decomposition into aleatoric and epistemic components. This appears to be consistent with the noise sensitivity displayed by Deep Ensemble method in Fig. 5.

The original loss function proposed in [29] did not work in this particular context of dose prediction. We found that regularizing the final layer with a sigmoid function and adding a mean-squared-error term led to a stable and effective implementation of Deep Evidential Learning. Upon completion of model training, we found that the model inherited uncertainty estimates at a more granular level compared to conventional MC Dropout and Deep Ensemble methods. Epistemic uncertainty ($U_{epis}$) in Deep Evidential Learning was highly correlated with prediction errors. Correlation indices (mutual information, Spearman’s coefficients) were comparable or higher than those for MC Dropout and Deep Ensemble methods (see Table 1). Aleatoric uncertainty $(U_{alea})$ demonstrated a more significant shift in its empirical CDF upon addition of Gaussian noise to CT intensity distribution compared to $U_{epis}$ (see Fig. 5). It also displayed visibly weaker correlation with prediction errors relative to $U_{epis}$. These traits of $U_{alea},U_{epis}$ appeared to be supportive of their conventional interpretations as reflecting data noise and model-related uncertainties respectively.

Another aspect of error-uncertainty relationship that appeared to distinguish Deep Evidential Learning was that the median error varied with uncertainty threshold much more linearly for $U_{epis}$ as compared to MC Dropout and Deep Ensemble methods, indicative of a more uniformly calibrated sensitivity to model errors (see Fig. 3, 4). This is compatible with the fact that the uncertainty heatmaps generated by the model were found to be highly effective in identifying potential regions of model’s inaccuracies.

An essential plan evaluation tool for the radiation oncologist is the DVH that characterizes a dose distribution. We demonstrated how Eve’s law of total variance enables one to express the predictive variance in terms of $U_{alea},U_{epis}$, and used this result to construct confidence intervals for DVH (see Fig. 8). This technique can be harnessed by the radiotherapy treatment planning team to enable a more statistically informed interpretation of DVH in deep learning-guided treatment planning. In the Knowledge-Based-Planning pipeline, dose prediction models can be used as inputs to some dose mimicking model [7] that produces deliverable treatments via optimization of a set of objective functions [25, 26]. The probabilistic dose distributions generated by our models can serve as robust inputs to the objective functions of [25, 26] to generate details of deliverable radiotherapy treatments. In this regard, it would be interesting to explore how the uncertainty estimates generated by our models translate to those of related quantities in the final treatment plan, such as the positions and motion dynamics of multileaf collimators and other machine parameters. This would pave the way for completing an uncertainty-aware AI-guided pipeline for radiotherapy treatment planning.

Finally, we outline several limitations of our work together with corresponding suggestions for future directions. Although the simple 3D U-Net backbone architecture led to efficient and fast convergence, it would be interesting to study how more complicated network structures like the transformers-based models of [17, 18] perform after modifying their final layers so that there are four output heads corresponding to the parameters of $\mathcal{P}$ in Deep Evidential Learning.

As a simple explicit example, a U-Net-like transformer model (Swin-UNet) was proposed in [51] and [52] for segmentation tasks. To enable Deep Evidential Learning on such backbone models, one can append a final 4-channel pointwise convolution layer to these transformer block-based architectures, which will yield the four parameters of the prior normal-inverse-gamma distribution in eqn. (4) as model outputs. Model training should employ the maximum likelihood loss function of [29] or the refined version that we have proposed here in eqn. (10). More generally, this simple minimal prescription of appending a final layer of 4-channel pointwise convolution layer applies for other transformer-based vision models adapted for segmentation. If higher model complexity is needed, this final layer can be preceded by more multi-channel pointwise convolution layers, with other essential traits of transformer models left intact, e.g. multi-head self-attention modules, shifted window-based attention mechanism, etc. Details of these inner layers are unaffected and do not require additional modifications when encapsulating the transformer-based vision model within the framework of Deep Evidential Learning.

Enabling Deep Evidential Learning on other backbone model structures will provide a more extensive study of the degree to which a low MAE can be attained while not compromising the uncertainty-error correlation. Indeed it will be interesting to explore how our refined Deep Evidential Learning model works for other regression tasks beyond the context of dose prediction, such as traffic forecasting in telecommunication networks in [53, 54]. It would be interesting to pursue whether the variety of model architectures examined in [53, 54] can be extended to uncertainty-aware ones using our techniques.

The MC Dropout and Deep Ensemble methods we studied for comparison purpose admit more complex variants which, in principle, may also yield aleatoric and epistemic uncertainties [36, 49, 55]. It would be interesting to perform a similar analysis for them, although they would bring with them additional challenges (e.g. the use of Laplacian priors for MC Dropout setting as described in [36]; computational cost in picking the optimal ensemble parameters as explained in [55] where a novel automated approach for Deep Ensemble was proposed).

We had used the same regularization term $\mathcal{L}_{reg}=\left|\log\left(\frac{y}{1-y}\right)-\gamma\right|(2\nu+\alpha)$ in the loss function following [29], which is independent of the $\beta$ parameter of the prior distribution $\mathcal{P}$ in eqn. (4). It would be interesting to explore if there are other alternatives which express full dependence on all four parameters of $\mathcal{P}$. In an analogous framework for classification, Sensoy et al. in [30] showed that the regularization term measuring the Kullback-Leibler divergence from an uninformative prior (e.g. uniform distribution) worked well for classification problems, and notably, the authors of [29] had attempted to find the corresponding version for regression yet without success.

Fundamentally, our modeling of dose prediction was based on the simplifying assumption that the patient’s CT images and the oncologist’s specification of the target regions are sufficient inputs for predicting the pareto-optimal dose distribution. In reality, effectiveness of radiation is often sensitive to a more complex web of biological factors. Nonetheless, we hope that our work can be a good starting point towards precision radiation oncology. Towards this goal, a pertinent future direction would be to expand the set of model’s input features to include genetic, molecular and other unique clinical characteristics of each patient, beyond just using medical images.