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

Summary

As 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.

Strengths

1. 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. 2. The presentation is clear and easy to follow. 3. The empirical performance is great comparing to baselines.

Weaknesses

1. 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. 2. The experiment section has some ambiguous points (see limitations below).

Questions

1. In line 306, why you say `without harming the accuracy` instead of `while promoting the accuracy` with Dice performance increases during reward tuning? 2. 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. 3. What does `efficient sampling` mean in the lower part of Figure 1? How does your sample algorithm differs from simple mini-batch sampling? 4. Typo in line 324: safter => safer.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

1. The paper lacks theoretical provements for FGRM like the uncertainty bounds, which is significant in the uncertainty estimation context. 2. 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.

Summary

This 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.

Strengths

1. The paper is well-written, logically organized, and effectively explains the novel aspects of the proposed framework. 2. The method has been rigorously tested on two large safety-critical surgical scene segmentation datasets, demonstrating superior performance.

Weaknesses

1. The evaluation metrics Uncertainty error mutual information (MI), Pixel Ratio (PR), and Box Ratio (BR) can not be found in [18]. 2. 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."

Questions

See the weaknesses.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

While the paper is largely comprehensive, there are some areas where it falls short: 1. 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. 2. Metric Citation Issue: There seems to be a mistake in the metric citation, which raises concerns about the overall quality of the experimental setup.

Summary

This 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. The 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. To 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.

Strengths

* 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. * 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.

Weaknesses

* 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.

Questions

* 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. * It would be nice to see uncertainty maps of other models to compare. Is this possible? * 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 * 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.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

None

Summary

This 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.

Strengths

The 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: 1. The paper is easy to understand and easy to follow, making it accessible to a wide audience. 2. The authors clearly demonstrated the experimental results and effectively derived novel findings, supporting their claims. 3. The experimental results strongly support the effectiveness of the proposed method, highlighting its superiority over existing approaches.

Weaknesses

The 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).

Questions

1. 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.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

1. 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. 2. 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.

Summary

This 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. [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.]

Strengths

- 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). - 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.

Weaknesses

- **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. - **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?)? - **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. - **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.

Questions

- Why is this considered a reinforcement learning algorithm? - How was reward maximization performed without the aleatoric and epistemic uncertainties provided from the evidential learning task?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The 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.

Response

Thank you for the detailed rebuttal and the additional experiments on i) Cityscapes and ii) results for vanilla EDL. I think this is a strong paper and the additional work provided by the authors strengthens the paper. My original score of 8 stands but I would defend publication of this manuscript. One note, the authors state in the rebuttal `Reply: Thanks for the comment. We would like to draw your attention to Table 1 and Table 2 in the paper, where we have compared our method with MC dropout. The results show that our method outperforms MC dropout by a large margin across all the evaluation metrics.` However, learned heteroscedastic uncertainty (https://arxiv.org/abs/1703.04977)) is not the same as MC dropout. In that paper, the authors develop a branched network (akin to hard-parameter sharing in multi-task learning) where one branch predicts the mean and the other the variance on a pixel-wise basis. This is fed into the loss function. MC-dropout can be used in addition to this method to estimate both epistemic and aleatoric uncertainty whereby the total uncertainty is the sum.

Thank you

Thank you for the effort during the rebuttal, I will keep the positive assessment of the paper.

Thank You and Our Further Response

Thank you very much for taking time to read our response. We aim to address your remaining concerns as follows. *Detailed Discussion and Comparisons*: Our paper has compared with a range of SOTA approaches, including probabilistic method [11], model ensemble-based methods [22][23], deep deterministic methods [10][35], and evidence-based method [6]. We have also included a comparison with an auxiliary network-based method ConfidNet in our response. We acknowledge the importance of a detailed discussion for in-depth comparison, but space limits prevented us from including such analysis in our original submission. Following your suggestion, we hereby provide a detailed analysis of pros and cons for each type of method. We will include them into the final version to enhance the writing quality of our comparisons. The probabilistic method MC Dropout [11] is simple in implementation using dropout layers, but it depends on multiple forward runs for uncertainty estimation. The model ensemble-based method Deep Ensembles [23] combines the outputs of multiple models for reliable uncertainty estimation, but the training of multiple models has computation burdens. Layer Ensemble [22] attaches multiple heads to intermediate layers of a network, achieving efficient uncertainty estimation with a single forward pass. Deep deterministic methods [35][10] quantify uncertainty through geometrical or statistical properties of hidden features, providing accurate out-of-distribution uncertainty estimation, but their Bi-lipschitz regularization can be unstable in deeper models. The evidence-based method [6] distinguishes aleatoric and epistemic uncertainty, but due to the lack of evidence ground truth, the model can only be trained with observed one-hot labels, which may have the tendency to peak the second-order distribution. ConfidNet designs an auxiliary network to learn a novel confidence criterion, making it applicable to any pre-trained segmentation model, albeit requiring an additional network for confidence estimation. Despite the notable achievements in these methods, they share one common limitation, i.e., relying on models trained on task objectives without considering the uncertainty estimation metric during the learning process. The superior performance of our method can be attributed to the following key factors. Firstly, our method explicitly optimizes uncertainty estimation metrics via a reward function, thereby directly calibrating prediction risk and model confidence. Furthermore, our fine-grained parameter updates scheme enables the effective parameter exploration of the policy network. In addition, our incorporated evidential learning layer allows us to provide more informative estimates of different uncertainty types. As shown in Table 1&2 in our paper, our method also has the advantage of less inference time, a critical factor for real-time intra-operative healthcare applications. The limitations of our method include the need for an additional RL reward maximization process and the current use of two reward functions for different types of uncertainty. *Medical Imaging Comparison*: For our comparison with auxiliary network-based methods in our rebuttal, we focus on ConfidNet (TPAMI’21) because it is a more recent work than [18] (MICCAI’19). Given that ConfidNet and [18] share similar insights by leveraging an auxiliary network, we only added ConfidNet in our comparison due to limited rebuttal time. However, we understand your point that [18] is a more relevant work as it is tailored for medical image segmentation. We are implementing [18] now and will post the results asap. Meanwhile, we would like to bring your attention to Table 1&2 in our paper, where we have included a MICCAI 2022 paper with the method of Layer Ensemble [22] for comparison. It is exactly a recent SOTA uncertainty estimation method designed for medical image segmentation. Our proposed method outperforms this MICCAI'22 paper by a large margin. Hope this can help relieve your concern. *Evaluation Metrics References*: Thank you for your further comments on direct comparison on the references. We agree that it is helpful to improve the paper, and actively try to do it. For the comparison with [17], since their code is not released, we tried our best to re-implement their method based on the information provided in their paper. Unfortunately, it is unstable in learning the joint latent space of input images and corresponding segmentation maps. As a result, the obtained uncertainty estimation was unsatisfactory, even presenting lower performance than baseline. We acknowledge the importance of including more SOTA medical imaging methods into comparison, so we are currently conducting experiments on [18] and [R3] to make our experimental evaluation more comprehensive. We believe our comparison with various types of general methods as well as dedicated medical imaging methods will effectively demonstrate the effectiveness of our method.

Terminology questions remain, but I have raised my score.

Thanks to the authors for their thorough responses to my and other reviews. Thanks to the authors for highlighting the connection to reference [31] and clarifying their view of the reinforcement learning framework. The rebuttal states: “segmentation predictions are actions and the input image is the state” — to me, this still doesn’t seem like an RL framework, because the actions do not affect the state of the system. However, it seems to me that [31] is characterized by a similar issue, and reframing the contribution may make the paper much easier to understand by future readers. I do think this is a good paper that should be accepted — at this point we are just discussing a terminology difference. I will defer to the AC and other reviewers on this point and will raise my score. I am satisfied with the response to all other points in my initial review, and especially appreciate the additional tests on endoscopic submucosal dissection (ESD) surgical data and improved visualizations of the results. I have also been following the discussion with reviewer jxwz about comparison to other medical imaging segmentation baselines. While the original paper already contained a thorough empirical characterization of the results, I think the latest response in the rebuttal further strengthens the paper and I am satisfied with the response on that point.