SlowGuess/ACL_2020_OCR huggingface dataset:
Abhyuday Jagannatha<sup>1</sup>, Hong Yu<sup>1,2</sup>
<sup>1</sup>College of Information and Computer Sciences, University of Massachusetts Amherst
$^{2}$ Department of Computer Science, University of Massachusetts Lowell
{abhyuday, hongyu}@cs.umass.edu
We address the problem of calibrating prediction confidence for output entities of interest in natural language processing (NLP) applications. It is important that NLP applications such as named entity recognition and question answering produce calibrated confidence scores for their predictions, especially if the applications are to be deployed in a safety-critical domain such as healthcare. However, the output space of such structured prediction models is often too large to adapt binary or multi-class calibration methods directly. In this study, we propose a general calibration scheme for output entities of interest in neural network based structured prediction models. Our proposed method can be used with any binary class calibration scheme and a neural network model. Additionally, we show that our calibration method can also be used as an uncertainty-aware, entity-specific decoding step to improve the performance of the underlying model at no additional training cost or data requirements. We show that our method outperforms current calibration techniques for named-entity-recognition, part-of-speech and question answering. We also improve our model's performance from our decoding step across several tasks and benchmark datasets. Our method improves the calibration and model performance on out-of-domain test scenarios as well.
Several modern machine-learning based Natural Language Processing (NLP) systems can provide a confidence score with their output predictions. This score can be used as a measure of predictor confidence. A well-calibrated confidence score is a probability measure that is closely correlated with the likelihood of model output's correctness. As a result, NLP systems with calibrated confidence can predict when their predictions are likely to be
incorrect and therefore, should not be trusted. This property is necessary for the responsible deployment of NLP systems in safety-critical domains such as healthcare and finance. Calibration of predictors is a well-studied problem in machine learning (Guo et al., 2017; Platt, 2000); however, widely used methods in this domain are often defined as binary or multi-class problems(Naeini et al., 2015; Nguyen and O'Connor, 2015). The structured output schemes of NLP tasks such as information extraction (IE) (Sang and De Meulder, 2003) and extractive question answering (Rajpurkar et al., 2018) have an output space that is often too large for standard multi-class calibration schemes.
Formally, we study NLP models that provide conditional probabilities  $p_{\theta}(y|x)$  for a structured output  $y$  given input  $x$ . The output can be a label sequence in case of part-of-speech (POS) or named entity recognition (NER) tasks, or a span prediction in case of extractive question answering (QA) tasks, or a relation prediction in case of relation extraction task.  $p_{\theta}(y|x)$  can be used as a score of the model's confidence in its prediction. However,  $p_{\theta}(y|x)$  is often a poor estimate of model confidence for the output  $y$ . The output space of the model in sequence-labelling tasks is often large, and therefore  $p_{\theta}(y|x)$  for any output instance  $y$  will be small. For instance, in a sequence labelling task with  $C$  number of classes and a sequence length of  $L$ , the possible events in output space will be of the order of  $C^L$ . Additionally, recent efforts (Guo et al., 2017; Nguyen and O'Connor, 2015; Dong et al., 2018; Kumar and Sarawagi, 2019) at calibrating machine learning models have shown that they are poorly calibrated. Empirical results from Guo et al. (2017) show that techniques used in deep neural networks such as dropout and their large architecture size can negatively affect the calibration of their outputs in binary and multi-class classification tasks.
Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2078-2092 July 5 - 10, 2020. ©2020 Association for Computational Linguistics
Parallelly, large neural network architectures based on contextual embeddings (Devlin et al., 2018; Peters et al., 2018) have shown state-of-the-art performance across several NLP tasks (Andrew and Gao, 2007; Wang et al., 2019). They are being rapidly adopted for information extraction and other NLP tasks in safety-critical applications (Zhu et al., 2018; Sarabadani, 2019; Li et al., 2019; Lee et al., 2019). Studying the miss-calibration in such models and efficiently calibrating them is imperative for their safe deployment in the real world.
In this study, we demonstrate that neural network models show high calibration errors for NLP tasks such as POS, NER and QA. We extend the work by Kuleshov and Liang (2015) to define well-calibrated forecasters for output entities of interest in structured prediction of NLP tasks. We provide a novel calibration method that applies to a wide variety of NLP tasks and can be used to produce model confidences for specific output entities instead of the complete label sequence prediction. We provide a general scheme for designing manageable and relevant output spaces for such problems. We show that our methods lead to improved calibration performance on a variety of benchmark NLP datasets. Our method also leads to improved out-of-domain calibration performance as compared to the baseline, suggesting that our calibration methods can generalize well.
Lastly, we propose a procedure to use our calibrated confidence scores to re-score the predictions in our defined output event space. This procedure can be interpreted as a scheme to combine model uncertainty scores and entity-specific features with decoding methods like Viterbi. We show that this re-scoring leads to consistent improvement in model performance across several tasks at no additional training or data requirements.
2 Calibration framework for Structured Prediction NLP models
Structured Prediction refers to the task of predicting a structured output  $y = [y_{1},y_{2},\dots y_{L}]$  for an input  $x$ . In NLP, a wide array of tasks including parsing, information extraction, and extractive question answering fall within this category. Recent approaches towards solving such tasks are commonly based on neural networks that are trained by
minimizing the following objective :
$$
\mathcal {L} (\theta | \mathcal {D}) = - \sum_ {i = 0} ^ {| \mathcal {D} |} \log \left(p _ {\theta} \left(y ^ {(i)} \mid x ^ {(i)}\right)\right) + R (\theta) \tag {1}
$$
where  $\theta$  is the parameter vector of the neural network and  $R$  is the regularization penalty and  $\mathcal{D}$  is the dataset  $\{(y^{(i)},x^{(i)})\}_{i = 0}^{|D|}$ . The trained model  $p_{\theta}$  can then be used to produce the output  $\hat{y} = \operatorname*{argmax}_{y\in \mathcal{Y}}p_{\theta}(y|x)$ . Here, the corresponding model probability  $p_{\theta}(\hat{y} |x)$  is the uncalibrated confidence score.
In binary class classification, the output space  $\mathcal{V}$  is  $[0,1]$ . The confidence score for such classifiers can then be calibrated by training a forecaster  $F_{y}:[0,1] \rightarrow [0,1]$  which takes in the model confidence  $F_{y}(P_{\theta}(y|x))$  to produce a recalibrated score (Platt, 2000). A widely used method for binary class calibration is Platt scaling where  $F_{y}$  is a logistic regression model. Similar methods have also been defined for multi-class classification (Guo et al., 2017). However, extending this to structured prediction in NLP settings is non-trivial since the output space  $|\mathcal{V}|$  is often too large for us to calibrate the output probabilities of all events.
Calibration methods for binary/multi class classification has been widely studied in related literature (Brocker, 2009; Guo et al., 2017). Recent efforts at confidence modeling for NLP has focused on several tasks like co-reference, (Nguyen and O'Connor, 2015), semantic parsing (Dong et al., 2018) and neural machine translation (Kumar and Sarawagi, 2019).
2.3 Calibration in Structured Prediction
In this section, we define the calibration framework by Kuleshov and Liang (2015) in the context of structured prediction problems in NLP. The model  $p_{\theta}$  denotes the neural network that produces an conditional probability  $p_{\theta}(y|x)$  given an  $(x,y)$  tuple. In a multi/binary class setting, a function  $F_{y}$  is used to map the output  $p_{\theta}(y|x)$  to a calibrated confidence score for all  $y \in \mathcal{V}$ . In a structured prediction setting, since the cardinality of  $\mathcal{V}$  is usually large, we instead focus on the event of interest set  $\mathcal{I}(x)$ .  $\mathcal{I}(x)$  contains events of interest  $E$  that are defined using the output events relevant to the deployment requirements of a model. The event  $E$  is a subset of  $\mathcal{V}$ . There can be several different schemes to define  $\mathcal{I}(x)$ . In later sections, we
discuss related work on calibration that can be understood as applications of different  $\mathcal{I}(x)$  schemes. In this work, we define a general framework for constructing  $\mathcal{I}(x)$  for NLP tasks which allows us to maximize calibration performance on output entities of interest.
We define  $F_{y}(E,x,p_{\theta})$  to be a function, that takes the event  $E$ , the input feature  $x$  and  $p_{\theta}$  to produce a confidence score between [0, 1]. We refer to this calibration function as the forecaster and use  $F_{y}(E,x)$  as a shorthand since it is implicit that  $F_{y}$  depends on outputs of  $p_{\theta}$ . We would like to find the forecaster that minimizes the discrepancy between  $F_{y}(E,x)$  and  $\mathbb{P}(y \in E|x)$  for  $(x,y)$  sampled from  $\mathbb{P}(x,y)$  and  $E$  uniformly sampled from  $\mathcal{I}(x)$ .
A commonly used methodology for constructing a forecaster for  $p_{\theta}$  is to train it on a held-out dataset  $D_{dev}$ . A forecaster for a binary classifier is perfectly calibrated if
$$
\mathbb {P} (y = 1 \mid F _ {y} (x) = p) = p. \tag {2}
$$
It is trained on samples from  $\{(x,\mathbb{I}(y = 1):(x,y)\in D_{dev}\}$ . For our forecaster based on  $\mathcal{I}(x)$ , perfect calibration would imply that
$$
\mathbb {P} (y \in E | F _ {y} (x, E) = p) = p. \tag {3}
$$
The training data samples for our forecaster are  $\{(x,\mathbb{I}(y\in E):E\in \mathcal{I}(x),(x,y)\in D_{dev}\}$
2.4 Construction of Event of Interest set  $\mathcal{I}(x)$
The main contributions of this paper stem from our proposed schemes for constructing the aforementioned  $\mathcal{I}(x)$  sets for NLP applications.
Entities of Interest : In the interest of brevity, let us define "Entities of interest"  $\phi(x)$  as the set of all entity predictions that can be queried from  $p_{\theta}$  for a sample  $x$ . For instance, in the case of answer span prediction for QA, the  $\phi(x)$  may contain the MAP prediction of the best answer span (answer start and end indexes). In a parsing or sequence labeling task,  $\phi(x)$  may contain the top-k label sequences obtained from viterbi decoding. In a relation or named-entity extraction task,  $\phi(x)$  contains the relation or named entity span predictions respectively. Each entity  $s$  in  $\phi(x)$  corresponds to a event set  $E$  that is defined by all outputs in  $\mathcal{V}$  that contain the entity  $s$ .  $\mathcal{I}(x)$  contains set  $E$  for all entities in  $\phi(x)$ .
Positive Entities and Events : We are interested in providing a calibrated probability for  $y \in E$  corresponding to an  $s$  for all  $s$  in  $\phi(x)$ . Here  $y$  is
the correct label sequence for the input  $x$ . If  $y$  lies in the set  $E$  for an entity  $s$ , we refer to  $s$  as a positive entity and the event as a positive event. In the example of named entity recognition,  $s$  may refer to a predicted entity span,  $E$  refers to all possible sequences in  $\mathcal{V}$  that contain the predicted span. The corresponding event is positive if the correct label sequence  $y$  contains the span prediction  $s$ .
Schemes for construction of  $\mathcal{I}(x)$  : While constructing the set  $\phi (x)$  we should ensure that it is limited to a relatively small number of output entities, while still covering as many positive events in  $\mathcal{I}(x)$  as possible. To explain this consideration, let us take the example of a parsing task such as syntax or semantic parsing. Two possible schemes for defining  $\mathcal{I}(x)$  are :
1. Scheme 1:  $\phi(x)$  contains the MAP label sequence prediction.  $\mathcal{I}(x)$  contains the event corresponding to whether the label sequence  $y' = \operatorname{argmax}_y p_\theta(y|x)$  is correct.
2. Scheme 2:  $\phi(x)$  contains all possible label sequences.  $\mathcal{I}(x)$  contains a event corresponding to whether the label sequence  $y'$  is correct, for all  $y' \in \mathcal{V}$
Calibration of model confidence by Dong et al. (2018) can be viewed as Scheme 1, where the entity of interest is the MAP label sequence prediction. Whereas, using Platt Scaling in a one-vs-all setting for multi-class classification (Guo et al., 2017) can be seen as an implementation of Scheme 2 where the entity of interest is the presence of class label. As discussed in previous sections, Scheme 2 is too computationally expensive for our purposes due to large value of  $|\mathcal{Y}|$ . Scheme 1 is computationally cheaper, but it has lower coverage of positive events. For instance, a sequence labelling model with a  $60\%$  accuracy at sentence level means that only  $60\%$  of positive events are covered by the set corresponding to  $\arg\max_y p_\theta(y|x)$  predictions. In other words, only  $60\%$  of the correct outputs of model  $p_\theta$  will be used for constructing the forecaster. This can limit the positive events in  $\mathcal{I}(x)$ . Including the top- $k$  predictions in  $\phi(x)$  may increase the coverage of positive events and therefore increase the positive training data for the forecaster. The optimum choice of  $k$  involves a trade-off. A larger value of  $k$  implies broader coverage of positive events and more positive training data for the forecaster training. However, it may also lead to
Calibration	BERT	BERT+CRF	DistilBERT
Platt	15.90±.03	15.56±.23	12.30±.13
Calibrated Mean	2.55±.34	2.31±.35	2.02±.16
+Var	2.11±.32	2.55±.32	2.73±.40
Platt+top2	11.4±.07	14.21±.16	11.03±.31
Calibrated Mean+top2	2.94±.29	4.82±.15	3.61±.17
+Var+top2	2.17±.35	4.26±.10	2.43±.16
+Rank+top2	2.43±.30	2.43±.45	2.21±.09
+Rank+Var+top2	1.81±.12	2.29±.27	1.97±.14
Platt+top3	17.46±.13	18.11±.16	12.84±.37
+Rank+Var+top3	3.18±.12	3.71±.25	2.05±.06
Table 1: ECE percentages on Penn Treebank for different models and calibration methods. The results are for top-1 MAP predictions on the test data. ECE standard deviation is estimated by repeating the experiments for 5 repetitions. ECE for uncalibrated BERT, BERT+CRF model and DistilBERT is  $35.11\%$ ,  $33.72\%$  and  $28.06\%$  respectively. heuristic- $k$  is 2 for all +Rank+Var+topk forecasters. Full feature model +Rank+Var+topk,  $k = 3$  is also provided for completeness.
an unbalanced training dataset that is skewed in favour of negative training examples.
Task specific details about  $\phi (x)$  are provided in the later sections. For the purposes of this paper, top-  $k$  refers to the top  $k$  MAP sequence predictions, also referred to as  $\argmax(k)$ .
2.5 Forecaster Construction
Here we provide a summary of the steps involved in Forecaster construction. Remaining details are in the Appendix. We train the neural network model  $p_{\theta}$  on the training data split for a task and use the validation data for monitoring the loss and early stopping. After the training is complete, this validation data is re-purposed to create the forecaster training data. We use an MC-Dropout(Gal and Ghahramani, 2016) average of (n=10) samples to get a low variance estimate of logit outputs from the neural networks. This average is fed into the decoding step of the model  $p_{\theta}$  to obtain top-  $k$  label sequence predictions. We then collect the relevant entities in  $\phi(x)$ , along with the  $\mathbb{I}(y \in E)$  labels to form the training data for the forecaster. We use gradient boosted decision trees (Friedman, 2001) as our region-based (Dong et al., 2018; Kuleshov and Liang, 2015) forecaster model.
Choice of the hyperparameter  $k$ : We limit our choice of  $k$  to  $\{2,3\}$ . We train our forecasters on training data constructed through top-2 and top-3 extraction each. These two models are then evaluated on top-1 extraction training data, and the best value of  $k$  is used for evaluation on test. This heuristic for  $k$  selection is based on the fact that
the top-1 training data for a good predictor  $p_{\theta}$ , is a positive-event rich dataset. Therefore, this dataset can be used to reject a larger  $k$  if it leads to reduced performance on positive events. We refer to the value of  $k$  obtained from this heuristic as a heuristic-  $k$ .
2.6 Feature Construction for Calibration
We use three categories of features as inputs to our forecaster.
Model and Model Uncertainty based features contain the mean probability obtained by averaging over the marginal probability of the "entity of interest" obtained from 10 MC-dropout samples of  $p_{\theta}$ . Average of marginal probabilities acts as a reduced variance estimate of un-calibrated model confidence. Our experiments use the pre-trained contextual word embedding architectures as the backbone networks. We obtain MC-Dropout samples by enabling dropout sampling for all dropout layers of the networks. We also provide  $10^{th}$  and  $90^{th}$  percentile values from the MC-Dropout samples, to provide model uncertainty information to the forecaster. Since our forecaster training data contains entity predictions from top-  $k$  MAP predictions, we also include the rank  $k$  as a feature. We refer to these two features as "Var" and "Rank" in our models.
Entity of interest based features contain the length of the entity span if the output task is named entity. We only use this feature in the NER experiments and refer to it as "ln".
Data Uncertainty based features: Dong et al. (2018) propose the use of language modelling (LM)
Calibration	BERT	BERT+CRF	DistilBERT
Baseline	60.30±.12	62.31±.11	60.17±.08
+Rank+Var+top2	60.30±.23	62.31±.09	60.13±.11
+Rank+Var+top3	59.84±.16	61.06±.14	58.95±.08
Table 2: Micro-avg f-score for POS datasets using the baseline and our best proposed calibration method. The confidence score from the calibration method is used to re-rank the events  $E \in \mathcal{I}(s)$  and the top selection is chosen. Standard deviation is estimated by repeating the experiments for 5 repetitions. Baseline refers to MC-dropout averaged (sample-size=10) output from the model  $p_{\theta}$ . heuristic-k is 2 for +Rank+Var+topk forecasters.
and OOV-word-based features as a proxy for data uncertainty estimation. The use of word-pieces and large pre-training corpora in contextual word embedding models like BERT may affect the efficacy of LM based features. Nevertheless, we use LM perplexity (referred to as "lm") in the QA task to investigate its effectiveness as an indicator of the distributional shift in data. Essentially, our analysis focuses on LM perplexity as a proxy for distributional uncertainty (Malinin and Gales, 2018) in our out-of-domain experiments. The use of word-pieces in models like BERT reduces the negative effect of OOV words on model prediction. Therefore, we do not include OOV features in our experiments.
We use BERT-base (Devlin et al., 2018) and distilBERT (Sanh et al., 2019) network architecture for our experiments. Validation split for each dataset was used for early stopping BERT fine-tuning and as training data for forecaster training. POS and NER experiments are evaluated on Penn Treebank and CoNLL 2003 (Sang and De Meulder, 2003), MADE 1.0 (Jagannatha et al., 2019) respectively. QA experiments are evaluated on SQuAD1.1 (Rajpurkar et al., 2018) and EMRQA (Pampari et al., 2018) corpus. We also investigate the performance of our forecasters on an out-of-domain QA corpus constructed by applying EMRQA QA data generation scheme (Pampari et al., 2018) on the MADE 1.0 named entity and relations corpus. Details for these datasets are provided in their relevant sections.
We use the expected calibration error (ECE) metric defined by Naeini et al. (2015) with  $N = 20$  bins (Guo et al., 2017) to evaluate the calibration of our models. ECE is defined as an estimate of the expected difference between the model confidence and accuracy. ECE has been used in several related works (Guo et al., 2017; Maddox et al., 2019;
Kumar et al., 2018; Vaicenavicius et al., 2019) to estimate model calibration. We use Platt scaling as the baseline calibration model. It uses the length-normalized probability averaged across 10 MC-Dropout samples as the input. The lower variance and length invariance of this input feature make Platt Scaling a strong baseline. We also use a "Calibrated Mean" baseline using Gradient Boosted Decision Trees as our estimator with the same input feature as Platt.
3.1 Calibration for Part-of-Speech Tagging
Part-of-speech (POS) is a sequence labelling task where the input is a text sentence, and the output is a sequence of syntactic tags. We evaluate our method on the Penn Treebank dataset (Marcus et al., 1994). We can define either the token prediction or the complete sequence prediction as the entity of interest. Since using a token level entity of interest effectively reduces the calibration problem to that of calibrating a multi-class classifier, we instead study the case where the predicted label sequence of the entire sentence forms the entity of interest set. The event of interest set is defined by the events  $y = \mathrm{MAP}_k(x)$  which denote whether each top- $k$  sentence level MAP prediction is correct. We use three choices of  $p_{\theta}$  models, namely BERT, BERT-CRF and distilBERT. We use model uncertainty and rank based features for our POS experiments.
Table 1 shows the ECE values for our baseline, proposed and ablated models. The value of heuristic-  $k$  is 2 for all  $+ \text{Rank} + \text{Var} + \text{topk}$  forecasters across all PTB models. "topk" in Table 1 refers to forecasters trained with additional top-  $k$  predictions. Our methods outperform both baselines by a large margin. Both "Rank" and "Var" features help in improving model calibration. Inclusion of top-2 prediction sequences also improve the calibration performance significantly. Table 1 also shows the performance of our full feature model "+"  $\text{Rank} + \text{Var} + \text{topk}$  for the sub-optimal value of
Calibration	CoNLL (BERT)	MADE 1.0 (bioBERT)
Platt	2.00±.12	4.00±.07
Calibrated Mean	2.29±.33	3.07±.18
+Var	2.43±.36	3.05±.17
+Var+ln	2.24±.14	2.92±.24
Platt+top3	16.64±.48	2.14±.18
Calibrated Mean+top3	17.06±.50	2.22±.31
+Var+top3	17.10±.24	2.17±.39
+Rank+Var+top3	2.01±.33	2.34±.15
+Rank+Var+ln+top3	1.91±.29	2.12±.24
Table 3: ECE percentages for the two named entity datasets and calibration methods. The results are for all predicted named entity spans in top-1 MAP predictions on the test data. ECE standard deviation is estimated by repeating the experiments for 5 repetitions. ECE for uncalibrated span margins from BERT model is  $3.68\%$  and  $5.59\%$  for CoNLL and MADE 1.0 datasets. heuristic- $k$  is 3 for all +Rank+Var+top3 forecasters.
Calibration	CoNLL (BERT)	MADE 1.0 (bioBERT)
Baseline	89.45±.08	84.01±.11
+Rank+Var+top3	89.73±.12	84.33±.07
+Rank+Var+ln+top3	89.78±.10	84.34±.10
Table 4: Micro-avg f-score for NER datasets and our best proposed calibration method. The confidence score from the calibration method is used to re-rank the events  $E \in \mathcal{I}(s)$  and a confidence value of 0.5 is used as a cutoff. Standard deviation is estimated by repeating the experiments for 5 repetitions. Baseline refers to MC-dropout averaged (sample-size=10) output of model  $p_{\theta}$ . heuristic-k is 3 for all +Rank+Var+top3 forecasters.