SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hongzheng Yang $^{1*}$ , Cheng Chen $^{2*}$ , Yueyao Chen $^{1}$ , Markus Scheppach $^{3}$ , Hon Chi Yip $^{1}$ , Qi Dou $^{1\dagger}$
$^{2}$ Harvard Medical School & Massachusetts General Hospital
<sup>3</sup> University Hospital of Augsburg
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-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.
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 targeted 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.
*Equal contributions (hzyang22@cse.cuhk.edu.hk, cchen101@mgh.harvard.edu)
† Corresponding author (qidou@cuhk.edu.hk)
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 fine-grained reward maximization (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.
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
Figure 1: The overview of our proposed fine-grained reward maximization (FGRM) framework for uncertainty estimation. The policy network is initialized with the parameters of a segmentation backbone, which is pre-trained with evidential learning to generate estimation for aleatoric and epistemic uncertainty separately. The policy network is tuned with a designed uncertainty estimation reward function by using a fine-grained parameter update scheme.
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.
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  $\pi_{\hat{\theta}}$  with model parameters  $\hat{\theta}$ , which maps the input image  $x$  to the segmentation prediction and the prediction uncertainty. The learned segmentation model  $\pi_{\hat{\theta}}$  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  $\pi_{\phi}$  with the same architecture of a pre-trained segmentation model  $\pi_{\hat{\theta}}$ . The details of the segmentation model pre-training will be described in Sec. 3.3. The policy network's parameter  $\phi$  is initialized with  $\hat{\theta}$ , then tuned to maximize a reward function  $R(\mu, \hat{y}, y)$ , where  $\mu, \hat{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:
$$
\mathcal {J} (\phi) = \mathbb {E} _ {(x, y)} [ R (\mu , \hat {y}, y) - \beta \log \left(\frac {\pi_ {\phi} (x)}{\pi_ {\hat {\theta}} (x)}\right) ], \tag {1}
$$
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  $\phi$  deviating too much from the pre-trained segmentation model  $\hat{\theta}$ . The hyper-parameter  $\beta$  controls the strength of the KL penalty.
In this regard, the policy network  $\pi_{\phi}$  can be updated to maximize the objective  $\mathcal{I}(\phi)$ . In particular, we desire the parameter  $\phi$  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  $\mathcal{J}(\phi)$ , 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
$\nabla_{\phi}\mathbb{E}_{(x,y)}[R(\mu ,\hat{y},y)] = \mathbb{E}_{(x,y)}[R(\mu ,\hat{y},y)\nabla_{\phi}\log \pi_{\phi}(\mu ,\hat{y} |x)]$  in which the gradient of each parameter is weighted by the reward  $R(\mu ,\hat{y},y)$  uniformly. We propose that the update for model  $\pi_{\phi}$  should not be identical everywhere in the parameter space. This is because that the parameters of  $\pi_{\phi}$  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 fine-grained 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  $\pi_{\phi}$  need to keep relatively close to those of the pre-trained segmentation model  $\pi_{\hat{\theta}}$ , 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:
$$
\operatorname {D i a g} | \mathcal {F} | = \operatorname {D i a g} (\mathbf {1} \odot \nabla_ {\phi} \log (\pi_ {\phi}) \nabla_ {\phi} \log (\pi_ {\phi}) ^ {T}). \tag {2}
$$
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:
$$
\phi = \phi + \eta \left(\sum_ {s} \left(\left(\operatorname {D i a g} | \mathcal {F} |\right) R (\mu , \hat {y}, y) \nabla_ {\phi} \log \pi_ {\phi} \left(\mu_ {s}, \hat {y} _ {s} | x\right)\right) - \beta \nabla_ {\phi} \log \left(\frac {\pi_ {\phi} (x)}{\pi_ {\hat {\theta}} (x)}\right)\right), \tag {3}
$$
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(\mathbf{p}|\boldsymbol{\alpha})$ , where  $\mathbf{p}$  is a simplex  $S_K = \{\mathbf{p}|\sum_{k=1}^{K} p_k = 1, p_k \geq 0\}$  representing the categorical probabilities with  $K$  segmentation classes. In this regard, the probability density function of Dirichlet distribution is given by:
$$
D (\mathbf {p} | \boldsymbol {\alpha}) = \left\{ \begin{array}{l l} \frac {1}{B (\boldsymbol {\alpha})} \prod_ {i = 1} ^ {K} p _ {i} ^ {\alpha_ {i} - 1} & \text {i f} \mathbf {p} \in S _ {K}, \\ 0 & \text {o t h e r w i s e ,} \end{array} \right. \tag {4}
$$
where  $B(\alpha)$  is the multinomial beta function as a normalizing factor. During the training process, the model iteratively measures 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 epochs  $E_{s}$ , RL training epochs  $E_{r}$ , reward maximization learning rate  $\eta$ , reward function  $R(\mu, \hat{y}, y)$ .
Output: model with reliable uncertainty estimation results
Segmentation Backbone Pre-training:
1: Initialize segmentation model  $\pi_{\theta}$
2: for  $e = 0,1,\ldots ,{E}_{s}$  do
3: sample batch of training data  $(\mathbf{x},\mathbf{y})$  from the training set
4: calculate  $L_{\mathrm{pretrain}}(\theta)$ $\triangleright$  Eq.(5)
5: update parameters:  $\theta \gets \theta -\eta \nabla_{\theta}L_{\mathrm{pretrain}}(\theta)$
6: end for
7: return pretrained parameters  $\hat{\theta}$
RL Reward Maximization:
1: initialize policy network  $\pi_{\phi}$  with pre-trained parameters  $\hat{\theta}$
2: for  $e = 0,1,\ldots ,{E}_{r}$  do
3: sample batch of training data  $(\mathbf{x},\mathbf{y})$  from the validation set
4:  $\operatorname{Diag}|\mathcal{F}| = \operatorname{Diag}(\mathbf{1} \oslash \nabla_{\phi}\log (\pi_{\phi})\nabla_{\phi}\log (\pi_{\phi})^{T})$ $\triangleright$  Eq.(2)
5:  $\phi \gets \phi +\eta (\sum_{s}((\mathrm{Diag}|\mathcal{F}|)R(\mu ,\hat{y},y)\nabla_{\phi}\log \pi_{\phi}(\mu_{s},\hat{y}_{s}|x)) - \beta \nabla_{\phi}\log (\frac{\pi_{\phi}(x)}{\pi_{\hat{\theta}}(x)}))$ $\triangleright$  Eq.(3)
6: end for
7: return tuned policy network  $\pi_{\phi}$  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  $\alpha$  with the model output  $\pi_{\hat{\theta}}(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  $\alpha_{k} / \sum_{k}(\alpha_{k})$  and the epistemic uncertainty can be quantified as the summation of Dirichlet distribution parameter  $\sum_{k}(\alpha_{k})$ . For maximum likelihood training, we treat the  $D(\mathbf{p}|\boldsymbol{\alpha})$  as the prior over the log likelihood  $\log (\text{Categorical}(y|p))$  and calculate the Bayes maximum likelihood loss by integrating out  $\mathbf{p}$  as:
$$
L _ {\text {p r e t r a i n}} (\theta) = \int \left(\sum_ {i = 1} ^ {K} - y _ {i} \log \left(p _ {i}\right)\right) \frac {1}{B (\boldsymbol {\alpha})} \prod_ {i = 1} ^ {K} \left(p _ {i}\right) ^ {\alpha_ {i} - 1} d \mathbf {p}. \tag {5}
$$
The model optimized by  $L_{\mathrm{pretrain}}$  could provide a proper initialization  $\hat{\theta} \in \operatorname{argmin}L_{\mathrm{pretrain}}(\theta)$  for RL training. At inference time, we can get the Dirichlet parameters with one forward pass of  $\pi_{\hat{\theta}}$ , thus quantify the reward and uncertainty efficiently.