SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 softmax function. For the base model pre-training, we use the Adam optimizer, with learning rate initialized to 1e-4. We totally trained 10 epochs 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(\mathrm{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.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)
Table 1: Comparison with state-of-the-art methods on laparoscopic cholecystectomy dataset for safety-critical tissue segmentation. Results are reported with mean±std of three independent runs.
Method	In-distribution Calibration	OOD Inference	Runtime (ms) ↓
ECE ↓	MI ↑	DICE ↑	PR ↑	BR ↑
Segmentation Backbone	19.46±1.12	2.88±0.76	71.22±1.22	0.88±0.33	0.08±0.04	0.052
Deep Ensemble [26]	16.20±2.01	3.82±0.86	72.59±1.44	0.92±0.23	0.12±0.04	0.201
Layer Ensemble [25]	17.04±0.97	3.94±0.59	72.23±1.04	1.00±0.25	0.16±0.07	0.072
DUM [38]	17.95±1.31	4.31±0.69	70.68±0.64	1.20±0.47	0.27±0.03	0.055
LDU [12]	14.42±0.72	4.93±0.45	73.13±1.57	1.44±0.44	0.35±0.09	0.054
MC-dropout [13]	17.27±0.60	3.31±0.49	72.42±1.57	1.03±0.22	0.15±0.06	0.510
LHU [21]	17.14±1.95	3.41±0.81	72.57±1.94	1.16±0.37	0.12±0.04	0.054
NatPN [6]	11.74±0.66	4.79±0.41	74.16±0.97	1.56±0.33	0.31±0.03	0.052
ConfidNet [9]	16.76±1.19	4.58±0.70	71.32±1.24	1.41±0.66	0.32 ±0.13	0.057
Auxiliary feat. [20]	16.87±1.92	4.35±0.75	71.32±1.24	1.33±0.55	0.28±0.11	0.055
FGRM (Ours)	9.63±0.74	5.87±0.47	74.88±0.91	1.85±0.25	0.47±0.02	0.052
Table 2: Comparison with state-of-the-art methods on endoscopic submucosal dissection dataset for safety-critical tissue segmentation. Results are reported with mean±std of three independent runs.
Method	In-distribution Calibration	OOD Inference	Runtime (ms) ↓
ECE ↓	MI ↑	DICE ↑	PR ↑	BR ↑
Segmentation Backbone	17.76±0.31	2.08±0.19	84.12±0.43	0.82±0.22	0.10±0.03	0.046
Deep Ensemble [26]	15.22±0.44	3.21±0.22	85.62±0.62	0.96±0.17	0.15±0.04	0.192
Layer Ensemble [25]	15.78±0.37	3.29±0.12	85.23±0.23	1.09±0.11	0.13±0.06	0.060
DUM [38]	16.07±0.54	3.51±0.27	83.11±0.85	1.14±0.10	0.15±0.07	0.047
LDU [12]	14.12±0.58	3.66±0.53	84.85±0.50	1.33±0.12	0.24±0.05	0.046
MC-dropout [13]	16.65±0.42	3.11±0.40	83.24±0.77	1.05±0.04	0.18±0.03	0.417
LHU [21]	16.62±0.62	3.36±0.34	83.31±1.22	1.12±0.31	0.14±0.04	0.047
NatPN [6]	13.67±0.91	4.16±0.53	86.77±0.93	1.45±0.22	0.32±0.07	0.046
ConfidNet [9]	15.52±1.12	3.53±0.65	84.41±0.61	1.19±0.25	0.26 ± 0.10	0.048
Auxiliary feat. [20]	15.76±1.15	3.41±0.61	84.41±0.88	1.14±0.40	0.18 ± 0.08	0.047
FGRM (Ours)	10.42±0.88	4.72±0.25	87.23±0.54	1.78±0.19	0.54±0.03	0.046
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 \times 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].
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  $\beta$  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
Figure 2: a) Estimated uncertainty map using our method on endoscopic submucosal dissection (ESD) of stomach dataset (first two columns) and laparoscopic cholecystectomy dataset (last two columns) for in-distribution calibration; b) Estimated uncertainty map using our method on ESD of esophagus (first two columns) and ESD of rectum (last two columns) for different organs calibration.
Figure 3: a) Effect of the three key components in our method, i.e., reward maximization, fine-grained parameter update, and evidential learning; b) Change of uncertainty estimation performance with the strength of KL penalty.
with the fine-grained parameter update, the uncertainty estimation is more stable with respect to different strengths of KL penalty. Since the hyperparameter  $\beta$  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  $2\mathrm{k}$  steps. This shows that our reward maximization improves the uncertainty estimation quality, as well as the accuracy of segmentation predictions.
Figure 4: a) Scatter plot of pixel-wise uncertainty with prediction correctness. b) Progression of ECE and Dice curves during reward maximization process.
We have proposed a novel reward optimization-based uncertainty estimation framework for safety-critical 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.
This work was supported in part by Hong Kong Innovation and Technology Commission Project No. ITS/237/21FP, in part by Hong Kong Research Grants Council Project No. T45-401/22-N, and in part by Science, Technology and Innovation Commission of Shenzhen Municipality Project No. SGDX20220530111201008.
[1] Moloud Abdar et al. “A review of uncertainty quantification in deep learning: Techniques, applications and challenges”. In: Information Fusion 76 (2021), pp. 243–297.
[2] 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.
[3] 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.
[4] Charles Blundell et al. "Weight uncertainty in neural network". In: International conference on machine learning. PMLR. 2015, pp. 1613-1622.
[5] Bertrand Charpentier, Daniel Zügner, and Stephan Gunnemann. "Posterior network: Uncertainty estimation without ood samples via density-based pseudo-counts". In: Advances in Neural Information Processing Systems 33 (2020), pp. 1356-1367.