Calibrating Structured Output Predictors for Natural Language Processing (2020)

Abstract: 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 are often too large to directly adapt binary or multi-class calibration methods. 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 tagging and Question Answering systems. We also observe an improvement in model 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.

## introduction 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 @xcite ; however, widely used methods in this domain are often defined as binary or multi-class problems @xcite . the structured output schemes of nlp tasks such as information extraction (ie) @xcite and extractive question answering @xcite 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 θ (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 θ (y|x) can be used as a score of the model's confidence in its prediction. however, p θ (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 θ (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 @xcite @xcite at calibrating machine learning models have shown that they are poorly calibrated. empirical results from @xcite 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. parallelly, large neural network architectures based on contextual embeddings @xcite have shown state-of-theart performance across several nlp tasks @xcite . they are being rapidly adopted for information extraction and other nlp tasks in safety-critical applications @xcite @xcite . 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 @xcite to define wellcalibrated 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-ofdomain 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. in binary class classification, the output space y is [0, 1]. the confidence score for such classifiers can then be calibrated by training a forecaster f y : [0, 1] → [0, 1] which takes in the model confidence f y (p θ (y|x)) to produce a recalibrated score @xcite . 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 @xcite . however, extending this to structured prediction in nlp settings is non-trivial since the output space |y| is often too large for us to calibrate the output probabilities of all events. ## experiments and results we use bert-base @xcite and dis-tilbert @xcite 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 @xcite corpus. we also investigate the performance of our forecasters on an out-of-domain qa corpus constructed by applying emrqa qa data generation scheme @xcite 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 @xcite with @xmath0. ece is defined as an estimate of the expected difference between the model confidence and accuracy. ece has been used in several related works @xcite @xcite to estimate model calibration. we use platt scaling as the baseline calibration model. it uses the lengthnormalized 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. ## discussion our proposed methods outperform the baselines in most tasks and are almost as competitive in others. features and top-k samples: the inclusion of top-k features improve the performance in almost all tasks when the rank of the prediction is included. we see large increases in calibration error when the top-k prediction samples are included in forecaster training without including the rank information in tasks such as conll ner and made 1.0 qa. this may be because the @xmath1, 2, 3 predictions may have similar model confidence and uncertainty values. therefore a more discriminative signal such as rank is needed to prioritize them. for instance, the difference between probabilities of @xmath2. in a sentence of length 10 or more, this difference in probability when normalized by length would account to very small shifts in the overall model confidence score. therefore an additional input of rank k leads to a substantial gain in performance for all models in pos. our task-agnostic scheme of "rank+var+topk" based forecasters consistently outperform or stay competitive to other forecasting methods. however, results from task-specific features such as "lm" and "len" show that use of task-specific features can further reduce the calibration error. our domain shift experimental setup has the same set of questions in both in-domain and out-of-domain datasets. only the data distribution for the answer passage is different. however, we do not observe an improvement in out-of-domain performance by using "lm" feature. a more detailed analysis of task-specific features in qa with both data and question shifts is required. we leave further investigations of such schemes as our future work. choice of k is important : the optimal choice of k seems to be strongly dependent on the inherent properties of the tasks and its output event set. in all our experiments, for a specific task all +rank+var+topk forecasters exhibit consistent behaviours with respect to the choice of k. in pos experiments, heuristic-@xmath3. in all other tasks, heuristic-@xmath4. our heuristic-k models are the best performing models, suggesting that the heuristic described in section 2.5 may generalize to other tasks as well. re-scoring : we show that using our forecaster confidence to re-rank the entities of interest leads to a modest boost in model performance for the ner and qa tasks. in pos no appreciable gain or drop in performance was observed for @xmath5. we believe this may be due to the already high token level accuracy (above 97%) on penn treebank data. nevertheless, this suggests that our re-scoring does not lead to a degradation in model performance in cases where it is not effective. our forecaster re-scores the top-k entity confidence scores based on model uncertainty score and entity-level features such as entity lengths. intuitively, we want to prioritize predictions that have low uncertainty over high uncertainty predictions, if their uncalibrated confidence scores are similar. we provide an example of such re-ranking in figure 2 . it shows a named entity span predictions for the correct span "such". the model p θ produces two entity predictions "off-spinner such" and "such". the un-calibrated confidence score of "off-spinner such" is higher than "such", but the variance of its prediction is higher as well. therefore the +rank+var+ln+top3 re-ranks the second (and correct) prediction higher. it is important to note here that the variance of "off-spinner such" may be higher just because it involves two token predictions as compared to only one token prediction in "such". this along with the "ln" feature in +rank+var+ln+top3 may mean that the forecaster is also using length information along with uncertainty to make this prediction. however, we see similar improvements in qa tasks, where the "ln" feature is not used, and all entity predictions involve two predictions (span start and end index predictions). these results suggest that use of uncertainty features are useful in both calibration and re-ranking of predicted structured output entities. out-of-domain performance : our experiments testing the performance of calibrated qa systems on out-of-domain data suggest that our methods result in improved calibration on unseen data as well. additionally, our methods also lead to an improvement in system accuracy on out-ofdomain data, suggesting that the mapping learned by the forecaster model is not specific to a dataset. however, there is still a large gap between the calibration error for within domain and out-of-domain testing. this can be seen in the reliability plot shown in figure 1 . the number of samples in each bin are denoted by the radius of the scatter point. the calibrated models shown in the figure corresponds to "+rank+var+lm+top3' forecaster calibrated using both in-domain and out-of-domain validation datasets for forecaster training. we see that out-of-domain forecasters are over-confident and this behaviour is not mitigated by using datauncertainty aware features like "lm". this is likely due to a shift in model's prediction error when applied to a new dataset. re-calibration of the forecaster using a validation set from the out-of-domain data seems to bridge the gap. however, we can see that the sharpness @xcite of out-of-domain trained, in-domain calibrated model is much lower than that of in-domain trained, indomain calibrated one. additionally, a validation dataset is often not available in the real world. mitigating the loss in calibration and sharpness induced by out-of-domain evaluation is an important avenue for future research. uncertainty estimation : we use mc-dropout as a model (epistemic) uncertainty estimation method in our experiments. however, our method is not specific to mc-dropout, and is compatible with any method that can provide a predictive distribution over token level outputs. as a result any bayesian or ensemble based uncertainity estimation method @xcite @xcite can be used with our scheme. in this work, we do not investigate the use of aleatoric uncertainty for calibration. our use of language model features is aimed at accounting for distributional uncertainty instead of aleatoric uncertainty @xcite . investigating the use of different types of uncertainty for calibration remains as our future work. ## conclusion we show a new calibration and confidence based re-scoring scheme for structured output entities in nlp. we show that our calibration methods outperform competitive baselines on several nlp tasks. our task-agnostic methods can provide calibrated model outputs of specific entities instead of the entire label sequence prediction. we also show that our calibration method can provide improvements to the trained model's accuracy at no additional training or data cost. our method is compatible with modern nlp architectures like bert. lastly, we show that our calibration does not over-fit on in-domain data and is capable of generalizing the calibration to out-of-domain datasets.