Title: Uncertainty Estimation for Safety-critical Scene Segmentation via Fine-grained Reward Maximization
Venue: NeurIPS 2023
Accepted: yes
Keywords: uncertainty estimation

semantic segmentation

medical application
Abstract: Uncertainty estimation plays an important role for future reliable deployment of deep segmentation models in safety-critical scenarios such as medical applications. However, existing methods for uncertainty estimation have been limited by the lack of explicit guidance for calibrating the prediction risk and model confidence. In this work, we propose a novel fine-grained reward maximization (FGRM) framework, to address uncertainty estimation by directly utilizing an uncertainty metric related reward function with a reinforcement learning based model tuning algorithm. This would benefit the model uncertainty estimation with direct optimization guidance for model calibration. Specifically, our method designs a new uncertainty estimation reward function using the calibration metric, which is maximized to fine-tune an evidential learning pre-trained segmentation model for calibrating prediction risk. Importantly, we innovate an effective fine-grained parameter update scheme, which imposes fine-grained reward-weighting of each network parameter according to the parameter importance quantified by the fisher information matrix. To the best of our knowledge, this is the first work exploring reward optimization for model uncertainty estimation in safety-critical vision tasks. The effectiveness of our method is demonstrated on two large safety-critical surgical scene segmentation datasets under two different uncertainty estimation settings. With real-time one forward pass at inference, our method outperforms state-of-the-art methods by a clear margin on all the calibration metrics of uncertainty estimation, while maintaining a high task accuracy for the segmentation results. Code is available at https://github.com/med-air/FGRM.

Uncertainty Estimation for Safety-critical Scene Segmentation via Fine-grained Reward Maximization

Abstract

Uncertainty estimation plays an important role for future reliable deployment of deep segmentation models in safety-critical scenarios such as medical applications. However, existing methods for uncertainty estimation have been limited by the lack of explicit guidance for calibrating the prediction risk and model confidence. In this work, we propose a novel fine-grained reward maximization (FGRM) framework, to address uncertainty estimation by directly utilizing an uncertainty metric related reward function with a reinforcement learning based model tuning algorithm. This would benefit the model uncertainty estimation through direct optimization guidance for model calibration. Specifically, our method designs a new uncertainty estimation reward function using the calibration metric, which is maximized to fine-tune an evidential learning pre-trained segmentation model for calibrating prediction risk. Importantly, we innovate an effective fine-grained parameter update scheme, which imposes fine-grained reward-weighting of each network parameter according to the parameter importance quantified by the fisher information matrix. To the best of our knowledge, this is the first work exploring reward optimization for model uncertainty estimation in safety-critical vision tasks. The effectiveness of our method is demonstrated on two large safety-critical surgical scene segmentation datasets under two different uncertainty estimation settings. With real-time one forward pass at inference, our method outperforms state-ofthe-art methods by a clear margin on all the calibration metrics of uncertainty estimation, while maintaining a high task accuracy for the segmentation results. Code is available at https://github.com/med-air/FGRM .

1 Introduction

Reliable segmentation of safety-critical scene is an important task for a variety of real-world scenarios such as autonomous driving [2, 3, 7, 40] and medical applications [19, 27, 29, 43]. It also plays a fundamental role for higher-level cognitive assistance relying on machine learning, for example, safety-critical soft tissue segmentation is essential for decision-making support towards autonomous robotic surgery [35]. In this regard, tolerance on prediction risk is extremely low, and estimating model uncertainty of the predicted result is equally as important as achieving a high task accuracy itself, especially when the targetted objects exhibit ambiguous boundaries in the scene. Therefore, it is inseparable to provide an estimation of the model uncertainty accompanied with the deterministic network predictions that may introduce risks under safety-critical scenarios.

2

†

Although uncertainty estimation of deep learning models for image segmentation has been intensively studied [1, 14], almost all previous methods solely resort to task objectives (e.g., cross-entropy, IoU or Dice) to train the network, and then investigate different inference strategies on the network, such as sampling multiple predictions to compute variance [24, 25, 26, 32], implicitly estimating uncertainty through feature density [12, 38], and designing evidence-based or deterministic layers to predict uncertainty [5, 6, 36]. Despite effectiveness to a certain extent, these practices commonly suffer from the issue of the absence of the uncertainty estimation metric to be considered, neither in the model training stage nor during the estimation process. This leads to indirect calibration of model confidence and unclear prior, which would affect the quality of estimated uncertainty, especially in the presence of ambiguous regions or out-of-distribution cases. Recently, reinforcement learning (RL) for model tuning has emerged with new achievements, such as improving language models from human feedback [33] and enhancing vision models with task rewards [28, 34]. We are inspired that the uncertainty estimation task can also be tackled with a RL reward maximization paradigm by innovating an uncertainty metric related reward function. This would have advantages over previous methods, because the model can be effectively fine-tuned with explicit guidance from the model confidence for calibrating the prediction risks.

The key challenges for achieving reward optimization-based uncertainty estimation lie in how to design the reward, and more importantly, how to update the network parameters during reward maximization. According to the latest works on vision tasks with RL [34, 42], the policy gradient methods estimate reward-weighted gradients for parameter updates, and the reward function is not necessarily differentiable. However, in these studies, the reward is uniformly weighted for all the network parameters, which may not be optimal because the influence of each model parameter on the reward function could differ from one to another. In addition, it can be expected that optimizing a network for dense segmentation prediction with feedback from a single reward could be challenging. In these regards, this process demands a meticulously crafted mechanism for parameter updates to identify the optimal solution within a constrained exploration space. Therefore, we consider that a designed fine-grained parameter update is required for model tuning process.

In this paper, we propose a novel f ineg rained r eward m aximization (FGRM) framework for uncertainty estimation in safety-critical scene segmentation. Specifically, we design a reward function that is directly related to uncertainty estimation, such as the calibration metrics, and demonstrate that by maximizing the reward function with reinforcement learning algorithms, the segmentation model can be explicitly tuned and well-calibrated for both segmentation prediction and uncertainty estimation. Importantly, to deal with the crucial yet challenging parameter updates with reward function, we leverage fisher information matrix, which carries the importance information of each parameter to network outputs, and yields a new fine-grained reward-weighting scheme for gradients of each network parameter. Moreover, we leverage the evidential deep learning framework to pre-train segmentation backbone, which makes the two types of uncertainties, i.e., aleatoric and epistemic uncertainty, to be separately estimated for reward maximization. Our main contributions are summarized as follows:

We propose a new reward optimization paradigm for uncertainty estimation, which can explicitly provide guidance to calibrate prediction risk and model confidence by designing a reward function that is closely tied to the uncertainty estimation.

We design an effective fine-grained parameter update mechanism for reward maximization, which imposes fine-grained reward-weighting for each network parameter and enables more effective model tuning based on the reward function.

We conduct extensive experiments on two medical datasets of safety-critical applications, including laparoscopic cholecystectomy scene segmentation and endoscopic submucosal dissection scene segmentation. Our new method consistently outperforms different types of state-of-the-art uncertainty estimation methods across all evaluation metrics.

2 Related Work

Uncertainty estimation for image segmentation. For the image segmentation task, related works on uncertainty estimation can be categorized into probabilistic methods [4, 11, 13, 39], model ensemble [24, 25, 26, 32], evidence-based methods [5, 6, 36], deep deterministic methods [12, 38] and auxiliary network-based methods [9, 20]. In specific, probabilistic methods estimate model uncertainty via dropout [13] and conditional variational autoencoders [4, 11, 39]. The ensemble-

Segmentation backbone

𝜶 = [𝛼1, 𝛼2 , … , 𝛼𝐾]

evidence

aleatoric

𝑝𝑘 =

𝛼𝑘

𝛴𝑘𝛼𝑘

epistemic

𝛴𝑘𝛼𝑘

Uncertainty quantification

Training loss

Evidence:

𝜶

=[1.5,1.5,1.5]

𝜶

=[5.0,5.0,5.0]

𝜶

=[1.0,2.0,2.0]

𝜶

=[2.0,4.0,8.0]

Policy Network

Sampled outputs

Per-sample  reward

0.52

0.34

0.41

0.55

Reward value

Efficient Inference

ID calibration

OOD inference

Train

Test

Alignment

Fine-grained

reward maximization

Initialization

Segmentation Backbone Pretraining

via Evidential Learning

Fisher information matrix

RL reward

maximization

Policy gradient

fine-grained update

Evidence layer

Cprerain(0) =

-log(P:)）B(a)

II(p)-dp

i=1

K

1.P

90.5

D(pla)

=

I-1

0

if p ∈ SK

otherwise

E

1.8

9.5

1.5

H

12.5

3+t

1.6

1499

Ouputs

F=Vlog（π)Vlog（π）

(Diag|F|)R(μ,,)Vlogπ(μ,9)

R(μ,,y)Vlogπ(μ,x)

based methods hold the insight to train a number of models for uncertainty estimation, and use stochastic segmentation networks [24, 32] to model a joint distribution over the entire map to reduce computation burdens. The evidence-based methods incorporate evidential layers into the network to model the conjugate prior of softmax prediction [31], which aims to directly predict the uncertainty of segmentation results. The deep deterministic methods can have low computational cost by exploiting the implicit feature density when the model architectures satisfy the Bi-lipschitz regularization [12]. Auxiliary network-based method designs an auxiliary branch to learn the confidence criterion, such as the True Class Probability (TCP). It requires an additional network branch for confidence estimation, and is often limited by the class imbalance problem, with a huge number of high-confidence samples during training. Despite much success previously achieved, these methods all resort to the model trained on task objectives, without taking into account the uncertainty estimation metric during learning. Our proposed method differs from them significantly, by exploring reinforcement learning algorithm to update the base segmentation model w.r.t. to a designed uncertainty reward.

Reinforcement learning for vision tasks. Using reinforcement learning for computer vision tasks is not completely new [28]. To improve the accuracy of scene segmentation, early methods use RL algorithms to iteratively refine the segmentation output [30], sequentially attend to various regions of the image [8], and realize multi-step interactive segmentation [37]. Very recently, there are some works exploring how to tune a base vision model to optimize the task-specific reward on different tasks, where the reward was designed as the task-specific evaluation metrics [34]. For panoptic segmentation, they adopt the per-sample panoptic quality [22] as the reward. For object detection task, they tune the model to optimize the recall and mean average precision metric. For tasks where the reward/metric are difficult to express by hand-crafted formulations such as the visual editing task, people propose to harness human feedback as instructions for reinforcement learning from human feedback (RLHF) [42]. The method needs human annotators to rank different outputs of the base model and then train a reward model that reflects human preferences. All these works demonstrated impressive performance on their respective tasks, and imply a promising trend to tackle vision tasks with reinforcement learning based model tuning. To the best of our knowledge, there is still no related work on addressing model uncertainty estimation with such a new RL-based paradigm.

3 Proposed Method

An overview of our proposed method is illustrated in Figure.1. Our method aims to learn a scene segmentation model capable of delivering not only highly accurate predictions but also high-quality uncertainty estimation for safety-critical scenes. We consider two types of uncertainty, i.e., aleatoric uncertainty and epistemic uncertainty, with aleatoric uncertainty accounting for irreducible data uncertainty, such as ambiguous tissue boundaries in surgical images, while epistemic uncertainty resulting from the lack of sufficient knowledge about unseen data. Let X denote the input image space and Y denote the label mask space, without loss of generality, we construct a differentiable segmentation model π ˆ θ with model parameters ˆ θ , which maps the input image x to the segmentation prediction and the prediction uncertainty. The learned segmentation model π ˆ θ is required to yield the estimations of aleatoric uncertainty for in-distribution (ID) data and epistemic uncertainty for out-of-distribution (OOD) data during real-time model inference.

3.1 Maximizing Uncertainty Estimation Reward with Reinforcement Learning

For existing uncertainty estimation approaches, a segmentation network is typically optimized by maximum likelihood estimation (MLE) based on the task objective function such as cross-entropy or Dice loss. However, in such a model training process, there is no explicit guidance to calibrate the prediction risk and model confidence. The challenge roots in the fact that there is usually no ground truth annotation for learning prediction uncertainty, thus it is hard to formulate a suitable optimization objective. To tackle this problem, we propose to make use of the reinforcement learning algorithm which explicitly maximizes a reward function that reflects uncertainty estimation metric, to tune the pre-trained segmentation network to calibrate the model parameters for both segmentation prediction and uncertainty estimation.

Specifically, we propose a policy network π ϕ with the same architecture of a pre-trained segmentation model π ˆ θ . The details of the segmentation model pre-training will be described in Sec. 3.3. The policy network's parameter ϕ is initialized with ˆ θ , then tuned to maximize a reward function R ( µ, ˆ y, y ) , where µ, ˆ y, y denote the uncertainty estimation, segmentation prediction and ground truth, respectively. We propose to design the reward function as calibration metrics that can quantify how closely the uncertainty is aligned with the model prediction risk, including both the in-distribution prediction confidence miscalibration and over-confidence in the out-of-distribution samples at inference. For the confidence miscalibration risk in in-distribution samples, the reward can be designed to minimize the difference in expectation between prediction accuracy and estimated aleatoric uncertainty. For the overconfidence risk in out-of-distribution samples, the reward can be designed to quantify whether the model is able to estimate higher epistemic uncertainty to OOD samples compared to their ID counterparts. The detailed formulation of the reward function is provided in the supplementary material. With the designed reward function, the RL objective can be expressed as:

where the first term is the reward maximization and the second term is a pixel-wise Kullback-Leibler (KL) penalty. The goal of the KL term is to mitigate over-optimization towards the reward and prevent policy network parameters ϕ deviating too much from the pre-trained segmentation model ˆ θ . The hyper-parameter β controls the strength of the KL penalty.

In this regard, the policy network π ϕ can be updated to maximize the objective J ( ϕ ) . In particular, we desire the parameter ϕ update can be fine-grained, and start with an easy-to-optimize initialization, i.e., the update should be specific for different parameters and the base model pretraining could provide a reasonable initialization for RL training. Moreover, the base model should have the ability to distinguish between the different types of uncertainty and ensure efficient sampling for reward computation. In the following, we will describe the design of RL-based parameter tuning and the base model pre-training in detail.

3.2 Fine-grained Parameter Update

To maximize the objective J ( ϕ ) , the main challenge is how to estimate the gradient of the reward term in Eq. 1. According to the policy gradient theorem [18], the gradient can be estimated as

∇ ϕ E ( x,y ) [ R ( µ, ˆ y, y )] = E ( x,y ) [ R ( µ, ˆ y, y ) ∇ ϕ log π ϕ ( µ, ˆ y | x )] , in which the gradient of each parameter is weighted by the reward R ( µ, ˆ y, y ) uniformly. We propose that the update for model π ϕ should not be identical everywhere in the parameter space. This is because that the parameters of π ϕ may contribute differently to the uncertainty estimation as well as segmentation prediction, therefore they can influence the reward function differently. With this insight, the uniform update scheme may not enable effective exploration for maximizing the uncertainty estimation reward. Furthermore, the dealing task here is quite challenging since we are optimizing a network to conduct dense prediction with feedback from the function of a single reward. Given the constrained exploration, together with guidance from the sparse reward value, finding a solution that achieves uncertainty estimation aligned with model risk is difficult.

To address this challenge, we propose to tune model parameters with uncertainty reward in a finegrained manner. In other words, we aim to weigh the reward for each network parameter individually. Our hypothesis is that since the parameters of policy network π ϕ need to keep relatively close to those of the pre-trained segmentation model π ˆ θ , this constraint limits the exploration space during RL training. If we can assign larger update weights to those parameters that significantly influence the uncertainty reward within this constrained exploration space, it can more efficiently reach the optimal solution for maximizing rewards. Inspired by elastic weight consolidation scheme in continual learning [23], we resort to the fisher information matrix, which carries information about the importance of each parameter to the model outputs. We compute the fisher information matrix as the covariance of the gradient of log likelihood as follows:

The calculation for fisher information matrix requires only the first-order derivative which can be computation efficient for large models. Next, we use the diagonal components in fisher information matrix to identify which parameters are more important and re-weight the parameter update during the RL-based model tuning. The updating scheme becomes as:

where s denotes each pixel. In this equation, the gradient of log likelihood for each parameter is weighted by both the reward value and diagonal element of the fisher information matrix. Larger value in the elements indicates a greater update step for network parameters, and such a gradient is computed for each individual pixel. The sum of pixel-wise gradients will be used to update the parameters. In this way, the parameter update is weighted by the information it carries about the final prediction, enabling more effective model tuning in the constrained parameter space which cannot deviate too much from the pre-trained parameters.

3.3 Segmentation Backbone Pre-training via Evidential Learning

To achieve effective reward maximization of uncertainty estimation, pre-training of the segmentation backbone also plays an important role, because it provides the capability to distinguish different uncertainty types that need to be well aligned with different sources of prediction risks. However, in a standard segmentation model, the final layer is often a softmax layer, which can only provide a point estimate of the predicted class probabilities. Their probability values produced by the softmax layer mix the aleatoric and epistemic uncertainty together, making it unclear to differentiate between the two types of uncertainty.

To address this issue, we borrow the idea of evidential deep learning [5] to train the model for explicitly parameterizing the conjugate prior of the categorical distribution. The segmentation of input x is modeled as a Dirichlet distribution D ( p | α ) , where p is a simplex S K = { p | ∑ K k =1 p k = 1 , p k ≥ 0 } representing the categorical probabilities with K segmentation classes. In this regard, the probability density function of Dirichlet distribution is given by:

where B ( α ) is the multinomial beta function as a normalizing factor. During the training process, the model iteratively measure the amount of support from data, i.e. evidence, in favor of the sample to be

Algorithm 1 FGRM algorithm

Input : segmentation backbone pre-training epoches E s , RL training epoches E r , reward maximization learning rate η , reward function R ( µ, ˆ y, y ) .

Output : model with reliable uncertainty estimation results

Segmentation Backbone Pre-training:

1: Initialize segmentation model π θ

2: for e = 0 , 1 , . . . , E s do

3: sample batch of training data (x, y) from the training set

5: update parameters: θ ← θ -η ∇ θ L pretrain ( θ )

6: end for

7: return pretrained parameters ˆ θ

RL Reward Maximization:

1: initialize policy network π ϕ with pre-trained parameters ˆ θ

2: for e = 0 , 1 , . . . , E r do

3: sample batch of training data (x, y) from the validation set

6: end for

7: return tuned policy network π ϕ with reliable uncertainty estimation

classified into a certain class. By replacing the last softmax layer with a non-negative evidence layer, we can represent the Dirichlet parameters α with the model output π ˆ θ ( x ) + 1 . In this way, the model is parameterized a distribution over the output categorical distribution. The aleatoric uncertainty can be quantified as expected categorical probability value α k / ∑ k ( α k ) and the epistemic uncertainty can be quantified as the summation of Dirichlet distribution parameter ∑ k ( α k ) . For maximum likelihood training, we treat the D ( p | α ) as the prior over the log likelihood log ( Categorical ( y | p )) and calculate the Bayes maximum likelihood loss by integrating out p as:

The model optimized by L pretrain could provide a proper initialization ˆ θ ∈ argmin L pretrain ( θ ) for RL training. At inference time, we can get the Dirichlet parameters with one forward pass of π ˆ θ , thus quantify the reward and uncertainty efficiently.

3.4 Implementation Details

Overall, the algorithm of our entire framework is presented in Algorithm 1. In our implementation, we employ an adapted TransUNet as segmentation backbone. We replace the last softmax layer with a non-negative evidence layer. The evidence layer is implemented by the softplus function. For the base model pre-training, we use the Adam optimizer, with learning rate initialized to 1e-4. We totally trained 10 epoches on the training set, with batch size 4. For the maximization of uncertainty estimation reward, we tune the base model to maximize the reward on a held-out validation set. For the ID miscalibration risk, we design the reward as the negative logarithm ECE function [15] -ln ( ECE ) . For the OOD over-confidence risk, we generate a corrupted version of the original validation and test set following [12, 16]. The reward was calculated as the uncertainty values ratio of the corrupted OOD samples compared to their ID counterparts. The model was updated by the policy gradient with our proposed fine-grained parameter update scheme. The learning rate and batch size was initialized as 1e-4 and 4, respectively.

4 Experiments

4.1 Experimental Datasets and Evaluation Metrics

We validate our method on two datasets of safety-critical surgical scene segmentations, i.e., laparoscopic cholecystectomy (LC) scene segmentation and endoscopic submucosal dissection (ESD)

scene segmentation. Dataset-1: For LC segmentation dataset, we adopt the public dataset CholecSeg8K [17], which contains 8,080 laparoscopic cholecystectomy image frames extracted from 17 video clips. Annotations are provided for segmentation of four different soft tissues, including abdominal wall, liver, gastrointestinal tract, and fat. Dataset-2: For ESD segmentation dataset, we collected a dataset with 1,203 image frames from 30 endoscopic surgical videos. The dataset was annotated by expert surgeons for the submucosal tissue, mucosa tissue, muscle tissue, and blood vessel. For data pre-processing, all the images were resized to 256 × 256 and normalized with zero mean and unit variance as the network inputs. For each dataset, we first randomly split 20% data for a held-out testing, and further split 80% of remaining data for training and 20% for validation. To demonstrate the broad applicability of our method, we also provide experimental evaluations on Cityscape [10] dataset for urban scene segmentation in Appendix A.5.

We evaluate our method under two uncertainty estimation scenarios, including in-distribution calibration and out-of-distribution inference. The inclusion of out-of-distribution inference aims to comprehensively assess the uncertainty estimation of our model when encountering data that deviates from their training distribution. The uncertainty estimation of in-distribution samples is conducted directly on the held-out test data, since there is no significant distribution shift between the training data and test data. The uncertainty estimation of out-of-distribution samples is conducted by generating a corrupted version of test data following [12, 16] and involving surgical video data of different organs. We adopt three commonly-used metrics for in-distribution calibration, including Expected Calibration Error (ECE) [15], Uncertainty-error mutual information (MI) [19] and Dice score, and two metrics for out-of-distribution inference, including the Pixel Ratio (PR) and Box Ratio (BR) [41].

4.2 Comparison with State-of-the-art Methods

Comparison state-of-the-art methods. Our method is compared with current state-of-the-art uncertainty estimation approaches, including ensemble-based methods of deep ensemble ( DE ) [26] and layer ensemble ( LE ) [25], deterministic uncertainty estimation methods of DUM [38] and LDU [12], probabilistic method of MC-dropout [13] and LHU [21], evidence-based method of NatPN [6] and auxiliary network based method of ConfidNet [9] and Auxiliary feat. [20]. In particular, DE simply uses an ensemble of neural networks trained independently to quantify the predictive uncertainty. LE lowers the computation cost of DE by ensembling from a single network's different depths. DUM trains a model with Bi-Lipschitz constraints and exploits the feature space densities to quantify the uncertainty. LDU advances scalable DUM by relaxing the Bi-Lipschitz constraints with a distinction maximization layer. MC-dropout approximates the Bayesian predictive uncertainty with a number of dropout runs. LHU designed the network consisted of two branches, which predicted the mean and variance separately. NatPN uses a normalizing flow to account for the epistemic uncertainty and proposes an input-dependent evidence update rule for the posterior distribution parameters. ConfidNet uses an additional network branch to learn the prediction confidence (True Class Probability) on a held-out validation set. Auxiliary feat. cascaded several consecutive convolution layers after the feature extractor network to learn to identify the segmentation errors. We also compare with a baseline model, which denotes training our Segmentation Backbone without using the evidential learning or any uncertainty estimation method.

Experimental results. The Table 1 shows the quantitative results on LC segmentation. We can see that different uncertainty estimation methods can generally improve the uncertainty estimation over baseline for both settings of the ID calibration and OOD inference. Our method achieves superior performance than the comparison methods on all the evaluation metrics and scenarios. Compared with the method NatPN, our approach obtains improvement of 2.1% in ECE and 0.29% in PR. Compared with MC-dropout that requires multiple forward runs to obtain uncertainty estimation, our method not only achieves higher uncertainty estimation quality, but also requires less inference time, which is important for real-time intra-operative healthcare applications. Table 2 shows the results on ESD segmentation. Similarly, compared to other methods, our approach achieves efficient uncertainty estimation and obtains higher performance on all evaluation metrics, with improvement of 3 . 25% on ECE and 0 . 83% on PR. The superiority of our method is attributed to the explicit model tuning with uncertainty estimation reward and the fine-grained parameter updates scheme that enables efficient parameter exploration of policy network. The visual results in Fig. 2(a) show that our method is able to produce varying degrees of uncertainty estimation according to the potential prediction risk. Besides the in-distribution result visualization, we also conducted additional tests on ESD surgical data involving out-of-distribution organs distinct from those covered by training data. Specifically, we applied the model trained on ESD surgical data from the stomach, to assess on ESD surgical data from the esophagus and rectum. Figure 2(b) showcases the results of our uncertainty estimation for these OOD cases, demonstrating the effectiveness of our method in providing meaningful uncertainty estimates when confronted with different organ types.

4.3 Ablation Analysis of the Proposed Method

We conduct experiments to investigate several key properties in our method: i) contribution of each key component; ii) effect of fine-grained parameter update; iii) relation between estimated uncertainty and prediction correctness; and iv) progression of uncertainty estimation and segmentation prediction.

Contribution of each component. We first validate the three key components in our method, i.e., uncertainty estimation reward maximization, fine-grained parameter update, and evidential learning, by continuously adding each component one by one to the plain MLE model. When evidential learning is not employed, the maximum class probability generated by the softmax function is used as both the aleatoric and epistemic uncertainty. As shown in Fig. 3 (a), each component plays an important role in improving the quality of uncertainty estimation on the two datasets. This demonstrates the benefits of using reinforcement learning with fine-grained parameter update and the evidential learning-based pre-training on improving uncertainty estimation.

Effect of fine-grained parameter update. We analyze the effect of fine-grained parameter update by varying the strengths of KL penalty, that is the hyperparameter β in Eq. 3. As observed in Fig. 3 (b), with the proposed fine-grained parameter update scheme, our method consistently obtains better uncertainty estimation performance under different strengths of KL penalty. Also we can see that

Mask prediction

Uncertainty

Ground truth

Error map

-------

(a)

(b)

(a)

(b)

LC Segmentation

22

LC Segmentation

18

22

ESD Segmentation

14

20

16

20

工

14

13

18

F 14

18

宁

12

ECE(%)

12

16

16

ECE

12

14

10

MI

ECE

12

8

10

(%)

CE

E

14

12

10 -

ECE

MI

工

H

10

8

(%)

IV

11

10

14

ESD Segmentation

8

6

4

工

T

工

6

4

8

6

4

工

工

工

4

6

ECE(%)

13 -

12

w/o fine-grained update

w/ fine-grained update

2

2

2

2

11

0

0

0

Q

Plain MLE

maximization

+Reward + Fine-grained + Evidential

update

learning

Plain MLE

maximization

+ Reward + Fine-grained + Evidential

update

learning

10

1e-2

2e-2

Strength of KL penalty

3e-2

4e-2

5e-2

with the fine-grained parameter update, the uncertainty estimation is more stable with respect to different strengths of KL penalty. Since the hyperparameter β controls the constrains on parameter optimization space during RL tuning, these results demonstrate that the fine-grained parameter update scheme contributes to effectively finding a better solution in the constrained parameter space.

Relation between estimated uncertainty and prediction correctness. We show the estimated uncertainty with the prediction correctness in Fig. 4 (a). We can see that the 'Plain MLE" model produces low uncertainty values for both incorrect and correct predictions, showing that the model is miscalibrated. With our FGRM method, the estimated uncertainty values are clearly related to the prediction correctness, that is the model is able to give high uncertainty values for incorrect predictions and low uncertainty values for correct predictions.

Progression of uncertainty estimation and segmentation prediction. We plot the ECE and dice curve during the reward maximization process to see their progression. As shown in Fig. 4 (b), we observe the model uncertainty calibration and Dice performance quickly increase during the reward tuning on the validation set and stay stable around 2k steps. This shows that our reward maximization improves the uncertainty estimation quality, as well as the accuracy of segmentation predictions.

(a)

LC Segmentation

ESD Segmentation

LC Segmentation

1.0

1.0

method

10.5

85

FGRM (Ours)

Plain MLE

ECE(%)

10.0

Dice(%)

0.8

0.8

9.5

80

Dice

Model Uncertainty

9.0

ECE

75

0.6

8.5

ESD Segmentation

0.4

0.4

9

92.5

ECE(%)

90.0

Dice(%)

0.2

0.2

7

Dice

ECE

85.0

0.0

0.0

6

Incorrect

Correct

Incorrect

Correct

82.5

0

500

1000

1500

2000

2500

Training Iterations

(b)

5 Conclusion

We have proposed a novel reward optimization-based uncertainty estimation framework for safetycritical scene segmentation. By maximizing a uncertainty estimation reward with a meticulously crafted fine-grained parameter update scheme, our framework provides explicit guidance to calibrate the prediction risk and model confidence with efficient network tuning. The superior efficacy of our method is demonstrated on two safety-critical surgical scene segmentation tasks, with both improved uncertainty estimation quality and segmentation predictions.

Limitations and future work. Since the calibration metrics for the uncertainty estimation in indistribution samples and out-of-distribution samples are different, our framework currently designed different reward functions to deal with the two scenarios. In the future, we would like to explore how to design a common reward function for model tuning such that a unified model can deliver high-quality uncertainty estimation under different situations. Additionally, we are planning to delve into more advanced RL techniques in order to reduce the frequency of reward calls and model updates. This will further enhance the training efficiency of our method.

Social Impact. Our FGRM framework provides a promising solution for efficiently estimating the uncertainty of model predictions, which is crucial for various real-world applications such as medical diagnosis or interventions. By producing high-quality uncertainty estimation, we can ensure the safer deployment of predictive models in practical applications.

References

Moloud Abdar et al. 'A review of uncertainty quantification in deep learning: Techniques, applications and challenges'. In: Information Fusion 76 (2021), pp. 243-297.

Victor Besnier, David Picard, and Alexandre Briot. 'Learning uncertainty for safety-oriented semantic segmentation in autonomous driving'. In: 2021 IEEE International Conference on Image Processing (ICIP) . IEEE. 2021, pp. 3353-3357.

Hermann Blum et al. 'Fishyscapes: A benchmark for safe semantic segmentation in autonomous driving'. In: proceedings of the IEEE/CVF international conference on computer vision workshops . 2019, pp. 0-0.

Charles Blundell et al. 'Weight uncertainty in neural network'. In: International conference on machine learning . PMLR. 2015, pp. 1613-1622.

Bertrand Charpentier, Daniel Zügner, and Stephan Günnemann. 'Posterior network: Uncertainty estimation without ood samples via density-based pseudo-counts'. In: Advances in Neural Information Processing Systems 33 (2020), pp. 1356-1367.

Bertrand Charpentier et al. 'Natural Posterior Network: Deep Bayesian Predictive Uncertainty for Exponential Family Distributions'. In: International Conference on Learning Representations . 2022.

Jiwoong Choi et al. 'Gaussian yolov3: An accurate and fast object detector using localization uncertainty for autonomous driving'. In: Proceedings of the IEEE/CVF International Conference on Computer Vision . 2019, pp. 502-511.

Wen-Hsuan Chu and Kris M Kitani. 'Neural batch sampling with reinforcement learning for semi-supervised anomaly detection'. In: Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XXVI 16 . Springer. 2020, pp. 751-766.

Charles Corbiere et al. 'Confidence estimation via auxiliary models'. In: IEEE Transactions on Pattern Analysis and Machine Intelligence 44.10 (2021), pp. 6043-6055.

Marius Cordts et al. 'The cityscapes dataset for semantic urban scene understanding'. In: Proceedings of the IEEE conference on computer vision and pattern recognition . 2016, pp. 32133223.

Michael Dusenberry et al. 'Efficient and scalable bayesian neural nets with rank-1 factors'. In: International conference on machine learning . PMLR. 2020, pp. 2782-2792.

Gianni Franchi et al. 'Latent Discriminant deterministic Uncertainty'. In: Computer VisionECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XII . Springer. 2022, pp. 243-260.

Yarin Gal and Zoubin Ghahramani. 'Dropout as a bayesian approximation: Representing model uncertainty in deep learning'. In: international conference on machine learning . PMLR. 2016, pp. 1050-1059.

Jakob Gawlikowski et al. 'A survey of uncertainty in deep neural networks'. In: arXiv preprint arXiv:2107.03342 (2021).

Chuan Guo et al. 'On calibration of modern neural networks'. In: International conference on machine learning . PMLR. 2017, pp. 1321-1330.

Dan Hendrycks and Thomas Dietterich. 'Benchmarking Neural Network Robustness to Common Corruptions and Perturbations'. In: Proceedings of the International Conference on Learning Representations (2019).

W-Y Hong et al. 'Cholecseg8k: a semantic segmentation dataset for laparoscopic cholecystectomy based on cholec80'. In: arXiv preprint arXiv:2012.12453 (2020).

Chloe Ching-Yun Hsu, Celestine Mendler-Dünner, and Moritz Hardt. 'Revisiting design choices in proximal policy optimization'. In: arXiv preprint arXiv:2009.10897 (2020).

Thierry Judge et al. 'CRISP-Reliable Uncertainty Estimation for Medical Image Segmentation'. In: Medical Image Computing and Computer Assisted Intervention-MICCAI 2022: 25th International Conference, Singapore, September 18-22, 2022, Proceedings, Part VIII . Springer. 2022, pp. 492-502.

Alain Jungo and Mauricio Reyes. 'Assessing reliability and challenges of uncertainty estimations for medical image segmentation'. In: Medical Image Computing and Computer Assisted Intervention-MICCAI 2019: 22nd International Conference, Shenzhen, China, October 13-17, 2019, Proceedings, Part II 22 . Springer. 2019, pp. 48-56.

Alex Kendall and Yarin Gal. 'What uncertainties do we need in bayesian deep learning for computer vision?' In: Advances in neural information processing systems 30 (2017).

Alexander Kirillov et al. 'Panoptic segmentation'. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition . 2019, pp. 9404-9413.

James Kirkpatrick et al. 'Overcoming catastrophic forgetting in neural networks'. In: Proceedings of the national academy of sciences 114.13 (2017), pp. 3521-3526.

Simon Kohl et al. 'A probabilistic u-net for segmentation of ambiguous images'. In: Advances in neural information processing systems 31 (2018).

Kaisar Kushibar et al. 'Layer Ensembles: A Single-Pass Uncertainty Estimation in Deep Learning for Segmentation'. In: Medical Image Computing and Computer Assisted InterventionMICCAI 2022: 25th International Conference, Singapore, September 18-22, 2022, Proceedings, Part VIII . Springer. 2022, pp. 514-524.

Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. 'Simple and scalable predictive uncertainty estimation using deep ensembles'. In: Advances in neural information processing systems 30 (2017).

Max-Heinrich Laves et al. 'Well-calibrated regression uncertainty in medical imaging with deep learning'. In: Medical Imaging with Deep Learning . PMLR. 2020, pp. 393-412.

Ngan Le et al. 'Deep reinforcement learning in computer vision: a comprehensive survey'. In: Artificial Intelligence Review (2022), pp. 1-87.

Hao Li et al. 'Region-based evidential deep learning to quantify uncertainty and improve robustness of brain tumor segmentation'. In: Neural Computing and Applications (2022), pp. 1-15.

Xuan Liao et al. 'Iteratively-refined interactive 3D medical image segmentation with multiagent reinforcement learning'. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition . 2020, pp. 9394-9402.

Andrey Malinin and Mark Gales. 'Predictive uncertainty estimation via prior networks'. In: Advances in neural information processing systems 31 (2018).

Miguel Monteiro et al. 'Stochastic segmentation networks: Modelling spatially correlated aleatoric uncertainty'. In: Advances in Neural Information Processing Systems 33 (2020), pp. 12756-12767.

Long Ouyang et al. 'Training language models to follow instructions with human feedback'. In: Advances in Neural Information Processing Systems 35 (2022), pp. 27730-27744.

André Susano Pinto et al. 'Tuning computer vision models with task rewards'. In: arXiv preprint arXiv:2302.08242 (2023).

Hamed Saeidi et al. 'Autonomous robotic laparoscopic surgery for intestinal anastomosis'. In: Science robotics 7.62 (2022), eabj2908.

Murat Sensoy, Lance Kaplan, and Melih Kandemir. 'Evidential deep learning to quantify classification uncertainty'. In: Advances in neural information processing systems 31 (2018).

Zhiqiang Tian et al. 'Multi-step medical image segmentation based on reinforcement learning'. In: Journal of Ambient Intelligence and Humanized Computing (2020), pp. 1-12.

Joost Van Amersfoort et al. 'Uncertainty estimation using a single deep deterministic neural network'. In: International conference on machine learning . PMLR. 2020, pp. 9690-9700.

Andrew G Wilson and Pavel Izmailov. 'Bayesian deep learning and a probabilistic perspective of generalization'. In: Advances in neural information processing systems 33 (2020), pp. 46974708.

Jingda Wu, Zhiyu Huang, and Chen Lv. 'Uncertainty-aware model-based reinforcement learning: methodology and application in autonomous driving'. In: IEEE Transactions on Intelligent Vehicles (2022).

Kilian Zepf et al. 'Laplacian Segmentation Networks: Improved Epistemic Uncertainty from Spatial Aleatoric Uncertainty'. In: arXiv preprint arXiv:2303.13123 (2023).

Shu Zhang et al. 'HIVE: Harnessing Human Feedback for Instructional Visual Editing'. In: arXiv preprint arXiv:2303.09618 (2023).

Ke Zou et al. 'TBraTS: Trusted brain tumor segmentation'. In: Medical Image Computing and Computer Assisted Intervention-MICCAI 2022: 25th International Conference, Singapore, September 18-22, 2022, Proceedings, Part VIII . Springer. 2022, pp. 503-513.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAs existing approaches for uncertainty estimation have been limited by the guidance for calibrating the prediction risk and model confidence, the paper proposes a novel fine-grained reward maximization (FGRM) framework, which addresses uncertainty estimation by reinforcement learning based model tuning with an uncertainty metric related reward function. It adopts the fisher information matrix for capturing parameter importance, acting as weights for fine-grained updates. Besides, evidential pre-training is incorporated to distinguish between aleatoric and epistemic uncertainty. The experimental results on two surgical datasets show FGRM improves uncertainty estimation for both ID and OOD data while not harming original segmentation performance.\n\nSTRENGTHS:\n1. The paper is well motivated for the uncertainty estimation problem, with reasonable application of fisher information matrix and evidential learning. Though built upon previous works that uses RL for guided training, I think the fine-grained update mechanism is novel and the contributions are clear.\n2. The presentation is clear and easy to follow.\n3. The empirical performance is great comparing to baselines.\n\nWEAKNESSES:\n1. While enjoying superior empirical performance, it would be better to provide some theoretical insights for FGRM like the uncertainty bounds, which is significant in the uncertainty estimation context.\n2. The experiment section has some ambiguous points (see limitations below).\n\nQUESTIONS:\n1. In line 306, why you say `without harming the accuracy` instead of `while promoting the accuracy` with Dice performance increases during reward tuning?\n2. As you use evidential learning to distinguish between aleatoric and epistemic uncertainty, what is the principle behind how it benefits your uncertainty estimation? You haven't explained this clearly in the ablation study.\n3. What does `efficient sampling` mean in the lower part of Figure 1? How does your sample algorithm differs from simple mini-batch sampling?\n4. Typo in line 324: safter => safer.\n\nLIMITATIONS:\n1. The paper lacks theoretical provements for FGRM like the uncertainty bounds, which is significant in the uncertainty estimation context.\n2. In Figure 3 and Figure 4(b), the paper doesn't record FGRM's performance on the OOD counterpart, which is an important contribution claimed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a novel Fine-Grained Reward Maximization (FGRM) framework to improve uncertainty estimation in deep segmentation models for safety-critical applications. The approach uses a reinforcement learning-based model tuning paradigm to optimize and calibrate the model. The FGRM framework is the first to leverage reinforcement learning for uncertainty estimation in safety-critical vision tasks, demonstrating improved performance on two surgical scene segmentation datasets.\n\nSTRENGTHS:\n1. The paper is well-written, logically organized, and effectively explains the novel aspects of the proposed framework. \n2. The method has been rigorously tested on two large safety-critical surgical scene segmentation datasets, demonstrating superior performance.\n\nWEAKNESSES:\n1. The evaluation metrics Uncertainty error mutual information (MI),  Pixel Ratio (PR), and Box Ratio (BR) can not be found in  [18].\n2. In the related works and experiments, lack of discussion and comparison with the auxiliary network-based method, e.g. [18] and Corbiere, Charles, et al. \"Confidence estimation via auxiliary models.\"\n\nQUESTIONS:\nSee the weaknesses.\n\nLIMITATIONS:\nWhile the paper is largely comprehensive, there are some areas where it falls short:\n\n1. Comparison with Other Methods: The paper claims superiority over state-of-the-art methods, however, it lacks a direct comparison or in-depth discussion with recent popular general uncertainty estimation methods, such as \"Confidence estimation via auxiliary models\" or methods specifically tailored for medical image segmentation like [18]. Including these comparisons would validate their claims more convincingly.\n\n2. Metric Citation Issue: There seems to be a mistake in the metric citation, which raises concerns about the overall quality of the experimental setup.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a novel method for uncertainty estimation. A segmentation network is first pre-trained by considering a generative model where the segmentation of an input $x$ is drawn from a Dirichlet distribution, which enables MLE. \n\nThe main contribution of the paper is then the reinforcement learning (RL) algorithm proposed whereby a novel reward function is used to maximise uncertainty estimation. Lastly, the authors posit that not all parameters in the network should be updated accordingly and use ideas from EWC in continual learning to learn parameter-specific update rules.\n\nTo learn aleatoric and epistemic uncertainty quickly, the authors assume a generative model where the segmentation is drawn from a Dirichlet distribution, thus optimising the MLE through integration of the conjugate prior and the likelihood function. The aleatoric and epistemic (calibrated) uncertainties can consequently be obtained through the learned $\\alpha$ parameter.\n\nSTRENGTHS:\n* This is an excellent paper. It is well written, clear in its intentions. The methodology is clearly presented (Figure 1 is excellent), it is clear how the algorithm works and why each section of the method was developed. The results compared to various baselines help consolidate the strength of the paper.\n\n* There are many novelties in this paper, which in isolation might already be interesting contributions. Taken as a whole, the authors have presented a compelling piece of work. The idea to combine i) evidential learning (parameterising the segmentation as $s ~ D(p|\\alpha)$ for uncertainty ii) using EWC to learn parameter updates necessary in the RL algorithm for calibrating uncertainty and iii) the RL reward function  are interesting and novel. The results are subsequently impressive.\n\nWEAKNESSES:\n* I would have liked to have seen evidence of the calibration on a more diverse set of datasets to better show the applicability and performance of the algorithm. A diverse set such as Cityscapes (real scenes), a dataset such as BraTS (tumour segmentation from MRI scans) in addition to the surgical videos used in this paper would have strengthened the work.\n\nQUESTIONS:\n* Did the authors consider also papers such as those which propose learning heteroscedastic uncertainty with MC dropout as a baseline such as in Kendal et al.(https://arxiv.org/abs/1703.04977). This method is used quite significantly and it would have been nice to see how it performs in comparison.\n\n* It would be nice to see uncertainty maps of other models to compare. Is this possible?\n\n* It is not entirely clear how the RL algorithm mitigates confidence miscalibration and OOD over confidence. This is only really mentioned in passing in Section 3.4 in the implementation details for Equation 1\n\n* In the experiments, did the authors consider a model where only evidential learning was used as another baseline to compare against vanilla MLE? It would be nice to understand how strong of a baseline that is and whether this could be used in isolation.\n\nLIMITATIONS:\nNone", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper empirically studied uncertainty estimation in safety-critical scene segmentation. The authors employed reinforcement learning (RL) methodologies, including fine-grained reward maximization (FGRM) framework and fisher information matrix for parameter updates. Additionally, to calibrate prediction risk and model confidence, the authors proposed a new reward function closely tied to uncertainty estimation. Through the experiments, the authors concluded the novel findings: (1) their method outperformed different types of state-of-the-art uncertainty estimation methods across all evaluation metrics; (2) the fine-grained parameter update mechanism improved the effectiveness of model tuning based on the reward function. Furthermore, the experimental results demonstrated the superiority of the proposed method on two medical datasets of safety-critical applications, specifically laparoscopic cholecystectomy scene segmentation and endoscopic submucosal dissection scene segmentation.\n\nSTRENGTHS:\nThe reviewer significantly understands the significance and importance of the task proposed in this paper. The findings and methodologies can be significantly utilized in safety-critical scenarios such as medical applications. In addition, the use of reinforcement learning (RL) for uncertainty estimation in model tuning is a novel approach, and the proposed reward function specifically designed for uncertainty estimation is innovative. Furthermore, since the proposed method consistently outperformed state-of-the-art uncertainty estimation methods on evaluation metrics, the contribution of this paper is noteworthy. Additionally, the strengths of this paper can be summarized as follows:\n\n1. The paper is easy to understand and easy to follow, making it accessible to a wide audience.\n\n2. The authors clearly demonstrated the experimental results and effectively derived novel findings, supporting their claims.\n\n3. The experimental results strongly support the effectiveness of the proposed method, highlighting its superiority over existing approaches.\n\nWEAKNESSES:\nThe reviewer agrees that the task addressed in this paper has several strengths and is significantly relevant to the community, but this paper is not technically novel. The main reasons are listed below (See Limitations).\n\nQUESTIONS:\n1. The application scope is significantly limited (only applicable to endoscopy modality and only for medical imaging). It would be better if the proposed framework could be applied to a broader range of applications. It would be beneficial to have further discussions on this issue.\n\nLIMITATIONS:\n1. Limited novelty. The differences (novelty) compared to existing methods are not clearly demonstrated. In the rebuttal stage, the reviewer strongly hopes that the authors will emphasize the novelty of their proposed method and clearly highlight its differences from existing approaches, clarifying what is new and innovative.\n\n2. Additionally, in the manuscript, the authors mention the inclusion of out-of-distribution (OOD) data alongside in-distribution data. The reviewer suggests justifying the use of OOD data in the experiments to strengthen the logical flow of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a new uncertainty estimation method for medical imaging segmentation tasks. The method relies on a pre-trained segmentation network that uses evidential learning to produce the parameters of a Dirichlet distribution over the class probabilities, which can be translated to aleatoric and epistemic uncertainty estimates. Then, an uncertainty estimation network is trained by updating its parameters based on a custom reward function weighted by the Fisher information matrix for all of the parameters. Experiments on two medical image segmentation datasets show that the method outperforms baselines including ensemble-based methods, deterministic uncertainty estimation methods, dropout, and an evidence update method. Ablations show the value of each part of the method.\n\n[Note: I do not work on uncertainty estimation or reinforcement learning, so other reviewers may be more qualified than me to comment here, and it is definitely possible I have misunderstood some of the paper.]\n\nSTRENGTHS:\n- The paper contains an extremely thorough ablation study and baseline comparison which demonstrate the value of each component of the system and a clear benefit over baselines (to be clear, this is a very big strength).\n- The idea of incorporating the outputs of evidential learning in the segmentation output into the training of an uncertainty estimation network seems interesting and is to my knowledge (from a brief literature search -- I am not familiar with this field) novel.\n\nWEAKNESSES:\n- **I do not understand how this is a reinforcement learning strategy, as opposed to standard training of another neural network that predicts uncertainty.** In particular, the form of the “RL Reward Maximization” part of Algorithm 1 looks exactly like standard training where network parameters are being updated via a loss function, and as far as I can tell, there is no notion of an “agent” performing actions in sequential time steps. I still think there is a contribution here, because the empirical results are clearly improved over the baselines, and there may be a misunderstanding on my part because I am not a reinforcement learning researcher. However, I think the paper may need a significant rewrite to clarify it and fully meet the bar for acceptance: either the paper should be written with the contribution being essentially a novel loss function, or the exposition needs some clarification in framing this as a reinforcement learning problem, as I did spend a good amount of time trying to understand this and could not figure it out. I am happy to update my rating later based on subsequent discussion on this point.\n- **Some details of the ablation study are unclear.** It appears evidential learning was added after reward maximization. How was reward maximization performed without the aleatoric and epistemic uncertainties provided from the evidential learning task (i.e. was R computed as in Appendix A.1? If so, how?)?\n- **OOD tasks are not realistic.** I do not think the OOD perturbations induced by Hendrycks et al are a realistic approximation to the types of OOD examples that would be seen in a medical segmentation task, where one would be especially concerned with unusual anatomy/pathology, etc. I don’t consider this a dealbreaker for accepting this paper, however, as it certainly seems reasonable that these OOD perturbations help in formulating the reward function; I am just less able to draw conclusions from the “OOD Inference” columns in the results tables.\n- **Visual results are hard to interpret.** Figure 2 does not provide the ground truth segmentations, so it is hard to tell where the model is incorrect, especially for readers unfamiliar with these types of images. Figure 1 in Appendix A.4 does show ground truth segmentations, but it would be helpful to see an error map of the segmentations to compare with the uncertainty.\n\nQUESTIONS:\n- Why is this considered a reinforcement learning algorithm?\n- How was reward maximization performed without the aleatoric and epistemic uncertainties provided from the evidential learning task?\n\nLIMITATIONS:\nThe paper has a clear “Limitations” section which describes a desire to unify the framework for dealing with in-distribution and out-of-distribution examples. I would like to see some additional discussion of the fact that that the OOD evaluations are not fully representative of what might be encountered in medical image segmentation tasks.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-27"}