Improving Uncertainty-Error Correspondence in Deep Bayesian Medical Image Segmentation
Prerak Mody Nicolas F. Chaves-de-Plaza Chinmay Rao Eleftheria Astrenidou Mischa de Ridder Nienke Hoekstra Klaus Hildebrandt Marius Staring



Purification of Gaussian States by Photon Subtraction
    Mattia Walschaers
    September 9, 2024
=====================================================




§ ABSTRACT 
	Increased usage of automated tools like deep learning in medical image segmentation has alleviated the bottleneck of manual contouring. This has shifted manual labour to quality assessment (QA) of automated contours which involves detecting errors and correcting them. A potential solution to semi-automated QA is to use deep Bayesian uncertainty to recommend potentially erroneous regions, thus reducing time spent on error detection. Previous work has investigated the correspondence between uncertainty and error, however, no work has been done on improving the “utility" of Bayesian uncertainty maps such that it is only present in inaccurate regions and not in the accurate ones. Our work trains the FlipOut model with the Accuracy-vs-Uncertainty (AvU) loss which promotes uncertainty to be present only in inaccurate regions. We apply this method on datasets of two radiotherapy body sites, c.f. head-and-neck CT and prostate MR scans. Uncertainty heatmaps (i.e. predictive entropy) are evaluated against voxel inaccuracies using Receiver Operating Characteristic (ROC) and Precision-Recall (PR) curves. Numerical results show that when compared to the Bayesian baseline the proposed method successfully suppresses uncertainty for accurate voxels, with similar presence of uncertainty for inaccurate voxels. Code to reproduce experiments is available at  <https://github.com/prerakmody/bayesuncertainty-error-correspondence>. 




	Bayesian Deep Learning, Bayesian Uncertainty, Uncertainty-Error Correspondence, Uncertainty Calibration, Contour Quality Assessment, Model Calibration














§ INTRODUCTION





In recent years, deep learning models are being widely used in radiotherapy for the task of medical image segmentation. Although these models have been shown to accelerate clinical workflows <cit.>, they still commit contouring errors <cit.>. Thus, a thorough quality assessment (QA) needs to be conducted, which places a higher time and manpower requirement on clinical resources. This creates a barrier to the adoption of such deep learning models <cit.>. Moreover, it also creates an obstacle for adaptive radiotherapy (ART) workflows, which have been shown to improve a patient's post-radiation quality-of-life <cit.>. This obstacle arises due to ART's need of regular contour updates. Currently, commercial auto-contouring tools do not have the ability to assist with quick identification and rectification of potentially erroneous predictions <cit.>. 


Quality assessment (QA) of incorrect contours would require two steps – 1) error detection and 2) error correction <cit.>. Currently, errors are searched for by manual inspection and then rectified using existing contour editing tools. Error detection could be semi-automated by recommending either potentially erroneous slices of a 3D scan <cit.>, or by highlighting portions of the predicted contours <cit.> or blobs <cit.>. Upon detection of the erroneous region, the contours could be rectified using point or scribble-based techniques <cit.> in a manner that adjacent slices are also updated. For this work, we will focus on error detection.


Various approaches to error detection have suggested using Bayesian Deep Learning (BDL) and the uncertainty that it can produce as a method to capture potential errors in the predicted segmentation masks <cit.>. Although such works established the potential usage of uncertainty in the QA of predictions, it may not be sufficient in a clinical workflow that relies on pixel-wise uncertainty as a proxy for error detection. In our experiments with deep Bayesian models, we observed that the relationship between prediction errors and uncertainty is sub-optimal, and hence has low clinical “utility". Ideally, for semi-automated contour QA, the uncertainty should be present only in inaccurate regions and not in the accurate ones. At times, literature usually refers to this as uncertainty calibration <cit.>, but we find this term incorrect as historically, calibration is referred to in context of probabilities of a particular event <cit.>. Thus, we believe it is semantically incorrect to say uncertainty calibration and instead propose to use the term uncertainty-error correspondence.






To create a Bayesian model that is incentivized to produce uncertainty only in inaccurate regions, we use the Accuracy-vs-Uncertainty (AvU) metric <cit.> and its probabilistic loss version <cit.> during training of a UNet-based  Bayesian model <cit.>. This loss promotes the presence of both accurate-if-certain (n_ac) as well as inaccurate-if-uncertain (n_iu) voxels in the final prediction (fig:abstract_visual). With uncertainty present only around potentially inaccurate regions, one can achieve improved synergy between clinical experts and their deep learning tools during the QA stage. Our work is the first to use the AvU loss in a dense prediction task like medical image segmentation and also with datasets containing natural and not synthetic variations as was previously done <cit.>. This work extends our conference paper <cit.> with additional datasets, experiments and metrics. There, we adapt the original AvU loss by considering the full theoretical range of uncertainties in the loss, rather than one extracted from the validation dataset <cit.>. For our work we use the predictive entropy as an uncertainty metric <cit.>.


Several other approaches have been considered in context of uncertainty, for e.g. ensembles, test time augmentation (TTA) and model calibration. While ensembles of models have good segmentation performance <cit.>, they are parameter heavy. TTA <cit.> performs inference by modulating a models inputs, but does not perform additional training, so may be unable to transcend its limitations. Calibration techniques attempt to make predictions less overconfident <cit.>, however they do not explicitly align model errors with uncertainty. All the above methods are benchmarked on the truthfulness of their output probabilities (when compared against voxel accuracies) using metrics like expected calibration error (ECE). However, a model with lower ECE than its counterparts may not necessarily have higher uncertainty-error correspondence.





Finally, to evaluate calibrative and uncertainty-error correspondence metrics, one needs to compute the “true" inaccuracy map. Similar to our conference paper <cit.> and inspired by <cit.>, we classify inaccuracies of predicted voxel maps into two categories: “errors" and “failures" (see sec:appendix_a). Segmentation “errors" are those inaccuracies which are considered an artifact similar to inter-observer variation, a phenomenon common in medical image segmentation <cit.>. Thus, we consider these smaller inaccuracies to be accurate in our computations, under the assumption they do not require clinical intervention. In the context of contour QA, such voxels should ideally be certain. Hence, only the segmentation “failures" are a part of the “true" inaccuracy map used to calculate the calibrative and uncertainty-error correspondence metrics.


To summarize, our contributions are as follows:
 
    
  * For the purpose of semi-automated quality assessment of predicted contours, we aim to improve uncertainty-error correspondence (unc-err) in a Bayesian medical image segmentation setting, pioneering this in the context of radiation therapy. Specifically, we propose using the loss form of the Accuracy-vs-Uncertainty (AvU) metric while training a deep Bayesian segmentation model.
    
  * We compare our Bayesian model with the AvU loss against an ensemble of deterministic models, five approaches employing calibration-based losses and also test time augmentation. We also perform an architectural comparison by comparing models with Bayesian convolutions placed in either the middle layers or decoder layers of a deep segmentation model.
    
  * We benchmark unc-err of the segmentation models on both in- and out-of-distribution radiotherapy datasets containing head-and-neck CT and Prostate MR scans. Models are benchmarked on these datasets across discriminative, calibrative and uncertainty-error correspondence metrics.














§ RELATED WORKS




 §.§ Epistemic and aleatoric uncertainty

Recent years have seen an increase in work that utilizes probabilistic modeling in deep medical image segmentation. The goal has been to account for uncertainty due to noise in the dataset (aleatoric uncertainty) as well as in the limitations of the predictive models learning capabilities (epistemic uncertainty). Noise in medical image segmentation refers to factors like inter- and intra- annotator contour variation <cit.> due to factors such as poor contrast in medical scans. Works investigating aleatoric uncertainty model the contour diversity in a dataset by either placing Gaussian noise assumptions on their output <cit.> or by assuming a latent space in the hidden layers and training on datasets containing multiple annotations per scan <cit.>. A popular and easy-to-implement approach to model for aleatoric uncertainty is called test-time augmentation (TTA) <cit.>. Here, different transformations of the image are passed through a model, and the resulting outputs are combined to produce both an output and its associated uncertainty. 

In contrast to aleatoric uncertainty, epistemic uncertainty could be used to identify scans (or parts of the scan) that are very different from the training dataset. Here, the model is unable to make a proper interpolation from its existing knowledge. Methods such as ensembling <cit.> and Bayesian posterior inference (e.g., Monte-Carlo DropOut, Stochastic Variational Inference) <cit.> are common methods to model epistemic uncertainty in neural nets. While Bayesian modeling is a more mathematically motivated and hence, principled approach to estimating uncertainty, ensembles have been motivated by the empirically-proven concept of bootstrapping. In contrast to Bayesian models where the perturbation is modelled by placing distributions on weights, ensembles use either different model weight initializations, or different subsets of the training data. In Bayesian inference techniques, perturbations are introduced in the models activation or weight space. Dropout <cit.> and DropConnect <cit.> are popular techniques that apply the Bernoulli distribution on these spaces. Stochastic variational inference (SVI) is another type of weight space perturbation that usually assumes the more expressive Gaussian distribution on the weights. Bayes by Backprop <cit.> and its resource-efficient variant such as FlipOut <cit.> are examples of SVI. For our work, we consider approaches that are designed for both epistemic uncertainty (Ensembles and SVI models) as well as aleatoric uncertainty (TTA).




 §.§ Uncertainty use during training

Other works also use the uncertainty from a base segmentation network to automatically refine its output using a follow-up network. This refinement network can be graphical <cit.> or simply convolutional <cit.>. Uncertainty can also be used in an active learning scenario, either with <cit.> or without <cit.> interactive refinement. Shape-based features of uncertainty maps have also been shown to identify false positive predictions <cit.>. Similarly, we too use uncertainty in our training regime, but with the goal of promoting uncertainty only in those regions which are inaccurate, an objective not previously explored in medical image segmentation.




 §.§ Model calibration

In context of segmentation, model calibration error is inversely proportional to the alignment of a models output probabilities with its pixel-wise accuracy. Currently there is no proof that reduction in model calibration error leads to improved uncertainty-error correspondence. However, a weak link can be assumed since both are derived from a models output probabilities. It is well known that the probabilities of deterministic models trained on the cross entropy (CE) loss are not well calibrated <cit.>. This means that they are overconfident on incorrect predictions and hence fail silently. This, which is an undesirable trait in context of segmentation QA and needs to be resolved.

To abate this overconfidence issue, methods such as post-training model calibration (or temperature scaling) <cit.>, ensembles <cit.>, calibration-focused training losses <cit.> and calibration-focused training targets <cit.> have been shown to improve model calibration for deterministic models. Temperature scaling, a post-training model calibration technique, has been shown to perform poorly in out-of-domain (OOD) settings <cit.>, relies wholly on an additional validation dataset and/or needs explicit shape priors <cit.>. FinerLocal temperature scaling techniques have been proposed that calibrate on the image or pixel level <cit.>, however they are still conceptually similar to the base method and are hence plagued by the same concerns. Others <cit.> used a shape prior module for out-of-domain robustness, but they only introduced synthetic textural variations in their work. 

Another approach to model calibration is to regularize a model during train to promote uncertainty. For e.g. the ECP <cit.> technique explicitly adds the negative entropy to the training loss. Conversely, the Focal loss <cit.> achieves thisattempts to calibrate a model implicitly by assigning lower weights (during training) to more confident predictions. Other methods smooth the hard targets of the ground truth towards a uniform distribution in the limit. For e.g. Label Smoothing <cit.> does this by modifying modifies the class distribution of a pixel by calculating a weighted average (using parameter α) between the hard target and a uniform distribution. On the other hand, Spatial Varying Label Smoothing (SVLS) <cit.> modifies a pixel's class allocation by considering classes around it. To avoid excessively making the models predictions uniform, Margin-based Label Smoothing (MBLS) <cit.> reformulates the above approaches by showing that they essentially perform loss optimization where an equality constraint is applied on a pixels logits. MBLS attempts to achieve the best discriminative-calibrative trade-off by softening this equality constraint. They subtract the max logit of a pixel with its other logits and only penalize those logit distances that are greater than a predetermined margin. Others extend thisthe MBLS framework by further tuningeither learning class-specific weights for the equality constraint <cit.> or reformulating SVLS to a similar formulation similar to MBLS <cit.>. Although these methods attempt to make models less overconfident, they do not explicitly align a model's error to its uncertainty.

There also exist other approaches to model calibration for e.g.,  multi-task learning <cit.>, mixup augmentation <cit.> and shape priors <cit.>. Multi-task learning requires additional data that may not always be present, while mixup creates synthetic samples which are not representative of the real data distribution. Finally, shape priors may not be applicable to tumors with variable morphology.

Model calibration techniques are evaluated by metrics like Expected Calibration Error (ECE) and its variants <cit.>, however others have also proposed terms like Uncertainty-Calibration Error (UCE) <cit.>. While ECE evaluates the equivalency between accuracy and predicted probability, UCE compares inaccuracy and uncertainty. However, while it is semantically appropriate to expect an average probability of p (0 ≤ p ≤ 1) to give the same average accuracy (i.e., the mathematical formulation of ECE), the same is not appropriate for inaccuracy and uncertainty u (0 ≤ u ≤ 1). Hence, UCE is not applicable to our work.

To conclude, the issue with each of the aforementioned techniques for epistemic, aleatoric and calibrative modeling is that they do not explicitly train the model to develop an innate sense of potential errors on a given segmentation task. Given that this is the primary requirement from a contour QA perspective, these models may be unable to have good uncertainty-error correspondence.




§ METHODS



 §.§ Neural Architecture

We adopt the OrganNet2.5D neural net architecture <cit.> which is a standard encoder-decoder model connected by four middle layers. It contains both 2D and 3D convolutions in the encoder and decoder as well as hybrid dilated convolutions (HDC) in the middle. This network performs fewer pooling steps to avoid losing image resolution and instead uses HDC to expand the receptive field. To obtain uncertainty corresponding to the output, we add stochasticity to the deterministic convolutional operations by replacing them with Bayesian convolutions <cit.>. We experiment with replacing deterministic layers in both the HDC as well as the decoder layers to understand the effect of placement.

In a Bayesian model, a prior distribution is placed upon the weights and is then updated to a posterior distribution on the basis of the training data. During inference (eq:1_variation_inference), we sample from this posterior distribution p(W|D) to estimate the output distribution p(y|x,D) with x, y and W being the input, output and neural weight respectively:

    p(y|x,D) = E_W ∼p(W|D)[p(y|x,W)].

This work uses a Bayesian posterior estimation technique called stochastic variational inference, where instead of finding the true, albeit intractable posterior, it finds a distribution close to it. We chose FlipOut-based <cit.> convolutions which assume the distribution over the neural weights to be a Gaussian and are factorizable over each hidden layer. Pure variational approaches would need to sample from this distribution for each element of the mini-batch <cit.>. However, the FlipOut technique only samples once and multiplies that random sample with a Rademacher matrix, making the forward pass less computationally expensive.



 §.§ Training Objectives


In this section , we use a notation format, where capital letters denote arrays while non-capital letters denote scalar values. 



  §.§.§ Segmentation Objective

Upon being provided a 3D scan as input, our segmentation model predicts for each class c ∈ C, a 3D probability map P̂_̂ĉ of the same size. Each voxel i ∈ N has a predicted probability vector P̂^̂î containing values p̂_̂ĉ^̂î for each class that sum to 1 (due to softmax). To calculate the predicted class of each voxel ŷ^̂î, we do:

    ŷ^̂î   = *argmax_c ∈ Cp̂_̂ĉ^̂î.
 

To generate a training signal, the predicted probability vector P̂^̂î is compared to the corresponding one-hot vector Y^i in the gold standard 3D annotation mask. Y^i is composed of y^i_c ∈{0,1}. 
Inspired by <cit.>, we re-frame the binary cross-entropy loss (eq:3_cross_entropy), as penalizing both the foreground (y_c^i=1) and background ((1-y_c^i)=1) voxels of the probability maps of each class with a weight w_c:





        



        





    L_CE   = - 1/|C|( ∑_c ∈ C w_c [      
        ∑_i ∈ N(
            y_c^iln(p̂_c^i)
            + 
            (1-y_c^i) ln(1 - p̂_c^i)
        ) ] ).


Note, we do not utilize the DICE loss for training as it has been shown to have lower model calibration metrics <cit.>. Also, since the CE loss is susceptible to fail during a class-imbalance, we use its weighted version.



  §.§.§ Uncertainty Objective

In a Bayesian model, multiple forward passes (m ∈ M) are performed and the output 3D probability maps (̂P̂_̂ĉ)̂_̂m̂ of each pass are averaged to output P̂_̂ĉ (eq:1_variation_inference). Using P̂_̂ĉ, we can calculate a host of statistical measures like entropy, mutual information and variance. We chose entropy as it has been shown to capture both epistemic uncertainty, which we explicitly model in FlipOut layers, as well as aleatoric uncertainty, which is implicitly modeled due to training data <cit.>. We use the predicted class probability vector P̂^̂î for each voxel and calculate its (normalized) entropy u^i:


    u^i    = -1/ln (|C|)∑_c ∈ Cp̂_̂ĉ^̂îln (p̂_̂ĉ^̂î).


Since we have access to the gold standard annotation mask, each voxel has two properties: accuracy and uncertainty. Accuracy is determined by comparing the gold standard class y^i to the predicted class ŷ^̂î. We use this to classify them in four different categories represented by n_ac, n_au, n_ic and n_iu, where n stands for the total voxel count and a, i, u, c represent the accurate, inaccurate, uncertain and certain voxels. A visual representation of these terms can be seen in fig:abstract_visual. Here, a voxel is determined to be certain or uncertain on the basis of a chosen uncertainty threshold t ∈ T where the maximum value in T is the maximum theoretical uncertainty threshold <cit.>. The aforementioned four terms are the building blocks of the Accuracy-vs-Uncertainty (AvU) metric <cit.> as shown in eq:5_avu - eq:6_avu_terms and it has a range between [0,1]. A higher value indicates that uncertainty is present less in accurate regions and more in inaccurate regions, thus improving the “utility" of uncertainty as a proxy for error detection.


    AvU^t   = n_ac^t + n_iu^t/n_ac^t + n_au^t + n_ic^t + n_iu^t
    
        n_ac^t    = ∑_i ∈{y_i = ŷ_̂î &
     u_i ≤ t} 1,     
        n_au^t = ∑_i ∈{y_i = ŷ_̂î &
     u_i > t} 1 
    
        n_ic^t    = ∑_i ∈{y_i ≠ŷ_̂î &
     u_i ≤ t} 1,      
        n_iu^t = ∑_i ∈{y_i ≠ŷ_̂î &
     u_i > t} 1


To maximize AvU for a neural net, one can turn it into a loss metric to be minimized. As done in <cit.> for an image classification setting, we minimize its negative logarithm (eq:7_avu_diff) to improve mathematical stability of gradient descent. However, the AvU metric, as defined above, is not differentiable with respect to the neural net's weights. This is due to all its constituent terms being produced either due to thresholding or max operations which introduce discontinuities that disrupt gradient flows.. This is because the model's outputs are simply used to create a mask and hence no backpropagation can take place. The AvU metric is made differentiable by instead using the uncertainty u^i derived from P̂^̂î (eq:4_entropy), thus allowing for gradient flows. by instead using the maximum value from each voxel's predicted probability vector p^i = max(P^i). Also, a smooth non-linear operation i.e., tanh is used to constrain its value (eq:8_avu_diff_terms_accurate). The differentiable uncertainty term is multiplied by other scalar weighing terms c.f. the maximum probability (p̂^̂î = max(P̂^̂î)) and accuracy/inaccuracy mask for a voxel. All these operations together allow us to calculate proxy values for n_ac, n_au, n_ic and n_iu. In addition, rather than evaluating the loss at a single uncertainty threshold, we integrate over the theoretical range of the uncertainty metric. Thresholding is done by once again multiplying the uncertainty value with a binary mask. The benefits of this thresholding were shown in our conference paper <cit.>:


    L_AvU^t   = - ln(1 + n_au^t + n_ic^t/n_ac^t + n_iu^t), 
    
            L_AvU   = 1/T∑_t ∈ T L_AvU^t,

where

    n_ac^t    = ∑_i ∈{y_i = ŷ_̂î &
     u_i ≤ t}p̂^̂î· (1 - tanh(u^i)), 
    
                n_ic^t    = ∑_i ∈{y_i ≠ŷ_̂î &
     u_i ≤ t} (1-p̂^̂î) · (1-tanh(u^i)),    
                
                n_au^t    = ∑_i ∈{y_i = ŷ_̂î &
     u_i > t}p̂^̂î·tanh(u^i) 
    
                n_iu^t    = ∑_i ∈{y_i ≠ŷ_̂î &
     u_i > t} (1-p̂^̂î) ·tanh(u^i).



Finally, the total loss L combines the segmentation and uncertainty loss as:

    L =L_CE + α· L_AvU.




 §.§ Evaluation




  §.§.§ Discriminative and Calibration Evaluation

We evaluate all models on the DICE coefficient for discriminative performance. Calibration is evaluated using the Expected Calibration Error (ECE) <cit.>. Numerical results are compared with the Wilcoxon signed-ranked test where a p-value ≤ 0.05 is considered significant.



  §.§.§ Uncertainty Evaluation

As the model is trained on the Accuracy-vs-Uncertainty (AvU) metric, we calculate the AvU scores up to the maximum normalized uncertainty of the validation dataset. A curve with the AvU score on the y-axis and the uncertainty threshold on the x-axis is made and the area-under-the-curve (AUC) for each scan is calculated. AUC scores provide us with a summary of the model performance regardless of the uncertainty threshold, and hence we use it to compare all models.   

The AvU metric outputs a single scalar value for the whole scan and does not offer much insight into the differences in uncertainty coverage between the accurate and inaccurate regions. To abate this, we compare the probability of uncertainty in inaccurate regions  p(u|i) to the probability of uncertainty in accurate regions p(u|a). Let us plot p(u|i) and p(u|a) on the y-axis and x-axis of a graph respectively, and define n_iu, n_au, n_ac and n_ic, as the count of true positives, false positives, true negatives and false negatives respectively. Thus, p(u|i) is the true positive rate and p(u|a) is the false positive rate. Computing this at different uncertainty thresholds provides us with the Receiver Operating Characteristic (ROC) curve, which we call the uncertainty-ROC curve <cit.>. 

Given that ROC curves are biased in situations with class imbalances between positive (inaccurate voxels) and negative (accurate voxels) classes, we also compute the precision-recall curves <cit.>.Here, precision is the probability of inaccuracy given uncertainty p(i|u) and recall is the probability of uncertainty given inaccuracy p(u|i). Note, that the precision-recall curves do not make use of n_ac, which can be high in count for a well-performing model.

Finally, to calculate the calibrative and uncertainty-correspondence metrics, we need an inaccuracy map. We use an inaccuracy map based on the concept of segmentation “failures" and “errors" (sec:appendix_a). To do this, we perform a morphological opening operation using a fixed kernel size of (3,3,1).




§ EXPERIMENTS AND RESULTS





 §.§ Datasets



  §.§.§ Head-and-Neck CT

Our first dataset contained Head and Neck CT scans of patients from the RTOG 0522 clinical trial <cit.>. The annotated data, which had been collected from the MICCAI2015 Head and Neck Segmentation challenge, contained 33 CT scans for training, 5 for validation and 10 for testing <cit.>. We further expanded the test dataset with annotations of 8 patients belonging to the RTOG trial from the DeepMindTCIA dataset (DTCIA) <cit.>. This dataset included annotations for the mandible, parotid glands, submandibular glands and brainstem. Although there were annotations present for the optic organs, we ignored them for this analysis as they are smaller compared to other organs and require special architectural design choices. Since the train and test patients came from the same study, we considered this as an in-distribution dataset. We also tested our models on the STRUCTSeg (50 scans) dataset <cit.>, hereby shortened as STRSeg. While the RTOG dataset contained American patients, the STRSeg dataset was made up of Chinese patients and hence considered out-of-distribution (OOD) in context of the training data. The uncertainties of this dataset were evaluated to a value of 0.4 since that is the maximum empirical normalized entropy.



  §.§.§ Prostate MR

Our second dataset contained MR scans of the prostate for which we use the ProstateX repository <cit.> containing 66 scans as the training dataset. The Medical Decathlon (Prostate) dataset with 34 scans <cit.> and the PROMISE12 repository with 50 scans <cit.> served as our test dataset. The Medical Decathlon dataset (abbreviated as PrMedDec henceforth) contained scans from the same clinic as the ProstateX training dataset. We combined the Peripheral Zone (PZ) and Transition Zone (TZ) from the MedDec dataset into 1 segmentation mask. The PROMISE12 dataset (abbreviated as PR12) was chosen for testing since literature <cit.> has shown lower performance on it and hence it serves as a good candidate to evaluate the utility of uncertainty. This dataset is different from ProstateX due to the usage of an endo-rectal coil in many of its scans as well as the presence of gas pockets in the rectum and dark shadows due to the usage of older MR machines. Thus, although these datasets contained scans of the prostate region, there exists a substantial difference in their visual textures. The maximum empirical normalized entropy of this 2-class dataset is 1.0 and hence the uncertainty-error correspondence metrics were calculated till this value.



 §.§ Experimental Settings

We tested the Accuracy-vs-Uncertainty (AvU) loss on four datasets containing scans of different modalities and body sites. We trained 11 models: Det (deterministic), Det+AvU, Ensemble, Focal, LS (Label Smoothing), SVLS (Spatially Varying Label Smoothing), MbLS (Margin based Label Smoothing), ECP (Explicit Confidence Penalty), TTA (Test-Time Augmentation), Bayes and Bayes + AvU. As the names suggest, Bayes and Bayes + AvU are Bayesian versions of the deterministic OrganNet2.5D model <cit.>. The baseline Bayes model contained Bayesian convolutions in its middle layers and was trained using only the cross-entropy (CE) loss. The Bayes + AvU was trained using both the CE and Accuracy-vs-Uncertainty (AvU) loss. Two additional Bayesian models were trained which tests if the placement of the Bayesian layers had any effect: BayesH and BayesH + AvU. Here, BayesH refers to the Bayesian model with Bayesian layers in the head of the model (i.e the decoder). Results for these models can be found in sec:app_bayesh.

The Ensemble was made of M=5 deterministic models with different initializations <cit.>. For TTA, we applied Gaussian noise and random pixel removals for M=5 times each and then averaged their outputs. The hyperparameters of the other models were chosen on the basis of the best discriminative, calibrative and uncertainty-error correspondence metrics on the validation datasets (<ref>). For the calibration focused methods we used the following range of hyperparameters: Focal (γ=1,2,3), MbLS (m=8,10,20,30) for head-and-neck CT, MbLS (m=3,5,8,10) for prostate MR, LS (α=0.1, 0.05, 0.01), SVLS (γ=1,2,3) and ECP (λ=0.1, 1.0, 10.0, 100.0) for head-and-neck CT and ECP (λ=0.1, 1.0, 10.0, 100.0, 1000.0) for prostate MR. For the AvU loss, we evaluated weighting factors in the range [10,100,1000,10000] for the head-and-neck dataset, and [100,1000,10000] for the Prostate dataset.

We trained our models for 1000 epochs using the Adam optimizer with a fixed learning rate of 10^-3. The deterministic model contained ≈ 550K parameters and thus the Ensemble contained ≈ 2.75M parameters. Since the Bayesian models double the parameter count in their layers they incurred an additional parameter cost and ended up with a total of ≈ 900K parameters. 













 §.§ Results
 
In sec:results_hn and sec:results_pros we show discriminative (DICE), calibrative (ECE) and uncertainty-error correspondence metrics (ROC-AUC, PRC-AUC) for the two datasets.



  §.§.§ Head-and-neck CT





















Results in tab:results_table_hn showed that the AvU loss on the Bayes model significantly improved calibrative and uncertainty-error correspondence (unc-err) metrics for both in-distribution (ID) and out-of-distribution (OOD) datasets. The Bayes+AvU model also always performed better than the Det, calibration-focused and TTA models for unc-err metrics. Also, its ECE scores were in most cases better than calibration-focused models. However, there was no clear distinction between the performance of the Ensemble and Bayes+AvU model for ECE and unc-err metrics across both datasets. Also, the AvU loss did not benefit the unc-err metrics for the Det model, in  both datasets. Of all the calibration-focused models, LS had the lowest ECE and unc-err metrics, while the ECP model had the best unc-err metrics. When compared to Det, the TTA model improved calibrative and unc-err metrics for the OOD dataset, while maintaining it for the ID dataset.

Visually, the Bayes+AvU model was able to successfully suppress uncertainty in the true positive (TP) (Case 1/2 in fig:results_visual_hn_rtog) and true negative (TN) (Case 3 in fig:results_visual_hn_rtog) regions of the predicted contour. Moreover, it also showed uncertainty in false positive (FP) regions while also suppressing uncertainty in TP regions (Case 3 in fig:results_visual_hn_structseg). Calibrative models (e.g. Focal, LS, SVLS) tended to be quite uncertain in TP or TN regions, which may lead to additional QA time. Detailed descriptions are provided in sec:app_visualresults.











  §.§.§ Prostate MR








Similar to the head-and-neck CT dataset, the use of the AvU loss on the baseline Bayes model significantly improved its uncertainty-error correspondence (unc-err) while maintaining calibration performance (tab:results_table_pros). Moreover, it improved the DICE values such that its one of the most competitive amongst all models. Also, the Bayes+AvU had better performance in both unc-err and calibrative metrics when compared to the Det, calibration-focused and TTA models. When comparing to the Ensemble, the Bayes+AvU had similar DICE. While Bayes+AvU had better calibrative and unc-err performance in the in-distribution (ID) dataset, the Ensemble performed better in the out-of-distribution (OOD) setting. The AvU loss had no positive effect on the DICE and unc-err performance of the Det model in both the ID and OOD setting, however there was an increase in ECE.

Visual results show that the Bayes+AvU successfully suppresses uncertainty in the true negative (Case 1 in fig:results_visual_pros_prmeddec, Case 2 in fig:results_visual_pros_pr12) and true positive (Case 2 in fig:results_visual_pros_prmeddec) regions of the predicted contour. It also shows uncertainty in the false positive regions (Case 2 in fig:results_visual_pros_prmeddec, Case 1/3 in fig:results_visual_pros_pr12)  





















§ DISCUSSION

Although medical image segmentation using deep learning can now predict high quality contours which can be considered clinically acceptable, a manual quality assessment (QA) step is still required in a clinical setting. To truly make these models an integral part of clinical workflows, we need them to be able to express their uncertainty and for those uncertainties to be useful in a QA setting. To this end, we test 11 models which are either Bayesian, deterministic, calibration-focused or ensembled.





 §.§ Discriminative and Calibrative Performance

In context of DICE and ECE, the use of the AvU loss on the baseline Bayes model always showed results which have never statistically deteriorated. Moreover, the DICE results for the in-distribution (ID) head-and-neck dataset (RTOG) were on-par with existing state-of-the-art models (83.6 vs 84.7 for <cit.>). The same held for the ID Prostate dataset (PRMedDec) where results were better than advanced models (84.9 vs 83.0 for <cit.>). These results validate the use of our neural architecture <cit.>, and training strategy.  

Secondly, although the Ensemble model, in general, had better or equivalent DICE and ECE scores across all 4 datasets, it also required 3x more parameters than the Bayes+AvU model. Also, as expected, and due to 5x more parameters, the Ensemble model performed better than the Det model for DICE and ECE. 

Finally, in the regime of segmentation “failures" as the inaccuracy map, the calibrative methods did not generally have improved calibration performance when compared to the Det model. In theory, these models regularize the model's probabilities by making it more uncertain and hence avoid overconfidence. In practice however, this leads to the predicted contours being uncertain along their accurate boundaries, most evident in visual examples of the Focal and SVLS model (see fig:results_visual_hn and fig:results_visual_pros). Also, visual image characteristics in different regions of the scan that are similar to the segmented organs may cause these models to showcase uncertainty in those areas (for e.g. patches of uncertainty in Case 3 of fig:results_visual_hn_rtog).









 §.§ Uncertainty-Error Correspondence Performance

Although calibrative metrics are useful to compare the average truthfulness of a model's probabilities, they may not be relevant to real-world usage in a pixel-wise segmentation QA scenario. Considering a clinical workflow in which uncertainty can be used as a proxy for error-detection, we evaluate the correspondence between them. Results showed that across both in- and out-of-distribution datasets, the Bayes+AvU model has one of the highest uncertainty-error correspondence metrics. Similar trends were observed for the BayesH+AvU (sec:app_bayesh) model, however Bayes+AvU was better. We hypothesize that this is due to perturbations in the bottleneck of UNet-like models having a better understanding of semantic concepts (e.g., shape, size etc) than the decoder layers. However, the AvU loss did not offer benefit to the Det model on both datasets indicating that this loss may rely on the model to already exhibit some level of uncertainty.

An interesting case is shown in fig:results_visual_hn_structseg (Case 3) which showed uncertainty on the white blob (a vein) in the middle of the grey tissue of the organ. Many models showed uncertainty on the vein due to a difference in its texture from that of the organ. However, this information may be distracting to a clinician as they are using uncertainty for error detection. Given that there were no segmentation “failures”, our Bayes+AvU model successfully suppressed all uncertainties. In another case (fig:results_visual_hn_rtog - Case 3), we saw that for 3D segmentation, uncertainty is also 3D in nature. Our Bayes+AvU model had an error in the second slice and correctly showed uncertainty there. However, this uncertainty overflowed on the first slice and hence penalized the uncertainty-error correspondence metrics. Such results indicate that during contour QA, the clinician can potentially trust our AvU loss models more than other models as they are better indicative of potential errors. This reduces time wasted analyzing false positive regions (i.e., accurate but uncertain) and hence increases trust between an expert and deep learning-based contour QA tools. Also note that in general, the two-class prostate dataset visually showcased higher levels of uncertainty than the six-class head-and-neck dataset. 




As seen in tab:results_table_hn, tab:results_table_pros and fig:results_curves, there is no clear choice between the top two performing models i.e., Bayes+AvU and Ensemble for uncertainty-error correspondence. The visual results, however, indicate that the Ensemble model is more uncertain in accurate regions. Also, for all the datasets, the Det model has high AvU scores when compared to the Bayes+AvU model (sec:appendix_c). Here, it is important to consider that the AvU metric (eq:6_avu_terms) is essentially uncertainty accuracy, and thus, also comes with its own pitfalls. Given that all models had a DICE value which leads to more accurate terms and less inaccurate terms, the AvU metric got skewed due to the large count of n_ac terms. However, upon factoring the ROC and PRC curves, it becomes evident that the Det model is not the best performing for uncertainty-error correspondence. 

Finally, all calibration-focused methods - Focal, ECP, LS, SVLS and MBLS had ROC and PRC metrics lower than the baseline Bayes model indicating that training for model calibration may not necessarily translate to uncertainty outputs useful for error detection. 



 §.§ Future Work

In a radiotherapy setting, the goal is to maximize radiation to tumorous regions and minimize it for healthy organs. This goal is often not optimally achieved due to imperfect contours caused by time constraints and amorphous region-of-interest boundaries on medical scans. Thus, an extension of our work could evaluate the contouring corrections made by clinicians in response to uncertainty-proposed errors in context of the dose changes to the different regions of interest. Such an experiment can better evaluate the clinical utility of an uncertainty-driven error correction workflow. 




§ CONCLUSION

This work investigates the usage of the Accuracy-vs-Uncertainty (AvU) metric to improve clinical “utility" of deep Bayesian uncertainty as a proxy for error detection in segmentation settings. Experimental results indicate that using a differentiable AvU metric as an objective to train Bayesian segmentation models has a positive effect on uncertainty-error correspondence metrics. We show that our AvU-trained Bayesian models have equivalent or improved uncertainty-error correspondence metrics when compared to various calibrative and uncertainty-based methods. Given that our approach is a loss function, it can be used with other neural architectures capable of estimating uncertainty.  


Given that deep learning models have shown the capability of reaching near expert-level performance in medical image segmentation, one of the next steps in their evolution is evaluating their clinical utility. Our work shows progress on this using a uncertainty-driven loss in a Bayesian setting. We do this for two radiotherapy body-sites and modalities as well in an out-of-distribution setting. Our hope is that the community is inspired by our positive results to further contribute to human-centric approaches to deep learning-based modeling.










The research for this work was funded by Varian, a Siemens Healthineers Company, through the HollandPTC-Varian Consortium (grant id 2019022) and partly financed by the Surcharge for Top Consortia for Knowledge and Innovation (TKIs) from the Ministry of Economic Affairs and Climate, The Netherlands.






The work follows appropriate ethical standards in conducting research and writing the manuscript, following all applicable laws and regulations regarding treatment of animals or human subjects.




We declare we don't have conflicts of interest.













§ SEGMENTATION “FAILURES" AND “ERRORS"



Gindraft=false





§ WEIGHTAGE OF AVU LOSS


The table below show the weights used for the AvU loss which were finetuned on the validation datasets of the head-and-neck CT and prostate MR. The final weightage was chosen by identifying the inflection point at which the ROC-AUC and PRC-AUC drop precipitously. Given that the AvU loss is a log term, its values are inherently small (≤ 1.0). This is then added to the cross-entropy term, which is a sum of logs (Eqn (3)) over all the voxels (=N) and all the classes (=C). Thus, we used a balancing term in the range of 10^1 to 10^4.







§ HYPERPARAMETER SELECTION



In the tables shown below, we report results for different hyperparameters of different model classes. If the DICE of a hyperparameter is 10.0 points lower than the class maximum, we ignore it. We also ignore models with large drops in ECE or AvU-AUC when compared to models in its own class. To choose the best hyperparameter, it has to perform as the best in four out of the five metrics, else we chose the middlemost hyperparameter.






























§ VISUAL RESULTS


Visual results in fig:results_visual_hn and fig:results_visual_pros show pairs of consecutive CT/MR slices to better understand the 3D nature of the output uncertainty across all models. We show examples with both high and low DICE to investigate the presence and absence of uncertainty in different regions of the model prediction. 



 §.§ Head-And-Neck CT

The first two rows of fig:results_visual_hn_rtog and fig:results_visual_hn_structseg show the mandible (i.e. lower jaw bone) with only the Bayes+AvU model having overall low uncertainty in accurate regions and high uncertainty in (or close to) inaccurate regions. 

In the next set of rows for head-and-necks CTs, we observe the parotid gland, a salivary organ, with (fig:results_visual_hn_rtog - Case 2) and without (fig:results_visual_hn_structseg - Case 2, Case 3) a dental scattering issue. In both cases, while the Det model shows low uncertainty, the baseline Bayes model shows high uncertainty in accurate regions. Usage of the AvU loss lowers uncertainty in these regions, while still exhibiting uncertainty in the erroneous regions, for e.g. the medial (i.e. internal) portion of the organ in fig:results_visual_hn_rtog (Case 2). 

Moving on to our last case, we see the submandibular gland, another salivary gland in fig:results_visual_hn_rtog (Case 3). The Ensemble, Focal, SVLS and MBLS models all display high uncertainty in the core of the organ, which are also accurately predicted. On the other hand, the AvU loss minimizes the uncertainty and shows uncertainty in the erroneous region on the second slice.



 §.§ Prostate MR

For the prostate datasets, we see two cases with high DICE in fig:results_visual_pros_prmeddec (Case 1) and fig:results_visual_pros_pr12 (Case 2) where the use of the AvU loss reduces uncertainty for the baseline Bayes model. 

We also see cases with low DICE in fig:results_visual_pros_prmeddec (Case 2) and fig:results_visual_pros_pr12 (Case 1). Due to their low DICE all models display high uncertainty, but the Bayes+AvU model shows high overlap between its uncertain and erroneous regions. The same is also observed in fig:results_visual_pros_pr12 (Case 3). 

Finally, in fig:results_visual_pros_prmeddec (Case 3), we do not see any clear benefit of using the AvU loss on the Bayes model.




§ BAYESH MODEL






