Calibrating Structured Output Predictors for Natural Language Processing

Structured Prediction refers to the task of predicting a structured output $y=[y_{1},y_{2},...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 :

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}^{|\mathcal{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{Y}$ is $[0,1]$. The confidence score for such classifiers can then be calibrated by training a forecaster $F_{y}:[0,1]\to[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{Y}|$ 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 (Bröcker, 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).

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{Y}$. In a structured prediction setting, since the cardinality of $\mathcal{Y}$ 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{Y}$ . 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}(\textit{E},x)$ and $\mathbb{P}(y\in\textit{E}|x)$ for $(x,y)$ sampled from $\mathbb{P}(x,y)$ and E uniformly sampled from $\mathcal{I}(\textit{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

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

The training data samples for our forecaster are $\{(x,\mathbb{I}(y\in\textit{E}):\textit{E}\in\mathcal{I}(x),(x,y)\in D_{dev}\}$.

The main contributions of this paper stem from our proposed schemes for constructing the aformentioned $\mathcal{I}(x)$ sets for NLP applications.

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\textit{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.

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

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=\text{MAP}_{k}(x)}$ which denote whether each top-$k$ sentence level MAP prediction is correct. We use three choice 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 +Rank+Var+topk forecasters across all PTB models. “top$k$” 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 “+Rank+Var+top$k$” for the sub-optimal value of $k=3$. It has lower performance than $k=2$ across all models. Therefore for the subsequent experimental sections, we only report top$k$ calibration performance using the heuristic-$k$ value only.

We use the confidence predictions of our full-feature model +Rank+Var+top$k$ to re-rank the top-$k$ predictions in the test set. Table 2 shows the sentence-level (entity of interest) accuracy for our re-ranked top prediction and the original model prediction.

For Named Entity (NE) Recognition experiments, we use two NE annotated datasets, namely CoNLL 2003 and MADE 1.0. CoNLL 2003 consists of documents from the Reuters corpus annotated with named entities such as Person, Location etc. MADE 1.0 dataset is composed of electronic health records annotated with clinical named entities such as Medication, Indication and Adverse effects.

The entity of interest for NER is the named entity span prediction. We define $\phi(x)$ as predicted entity spans in $\text{argmax}(k)$ label sequences predictions for $x$. We use bert-base with token-level softmax output and marginal likelihood based training. The model uncertainty estimates for “Var” feature are computed by estimating the variance of length normalized MC-dropout samples of span marginals. Due to the similar trends in behavior of bert and bert+CRF model in POS experiments, we only use bert model for NER. However, the span marginal computation can be easily extended to linear-chain CRF models. We also use the length of the predicted named entity as the feature “ln” in this experiment. Complete details about forecaster and baselines are in the Appendix. Value of heuristic-$k$ is 3 for all +Rank+Var+topk forecasters. We show ablation and baseline results for $k=3$ only. However, no other forecasters for any $k\in\{2,3\}$ outperform our best forecasters in Table 3.

We use the confidence predictions of our “+Rank+Var+top3” models to re-score the confidence predictions for all spans predicted in top-$3$ MAP predictions for samples in the test set. A threshold of 0.5 was used to remove span predictions with low confidence scores. Table 4 shows the Named Entity level (entity of interest) Micro-F score for our re-ranked top prediction and the original model prediction. We see that re-ranked predictions from our models consistently improve the model f-score.

We use three datasets for evaluation of our calibration methods on the QA task. Our QA tasks are modeled as extractive QA methods with a single span answer predictions. We use three datasets to construct experiments for QA calibration. SQuAD1.1 and EMRQA (Pampari et al., 2018) are open-domain and clinical-domain QA datasets, respectively. We process the EMRQA dataset by restricting the passage length and removing unanswerable questions. We also design an out-of-domain evaluation of calibration using clinical QA datasets. We follow the guidelines from Pampari et al. (2018) to create a QA dataset from MADE 1.0 (Jagannatha et al., 2019). This allows us to have two QA datasets with common question forms, but different text distributions. In this experimental setup we can mimic the evaluation of calibration methods in a real-world scenario, where the task specifications may remain the same but the underlying text source changes. Details about dataset pre-processing and construction are provided in the Appendix.

The entity of interest for QA is the top-$k$ answer span predictions. We use the “lm” perplexity as a feature in this experiment to analyze its behaviour in out-of-domain evaluations. We use a $2$ layer unidirectional LSTM to train a next word language model on the EMRQA passages. This language model is then used to compute the perplexity of a sentence for the “lm” input feature to the forecaster. We use the same baselines as the previous two tasks.

Based on Table 5, our methods outperform the baselines by a large margin in both in-domain and out-of-domain experiments. Value of heuristic-$k$ is 3 for all +Rank+Var+topk forecasters. We show ablation and baseline results for $k=3$ only. However, no other forecasters for any $k\in\{2,3\}$ outperform our best forecasters in Table 5 . Our models are evaluated on SQuAD1.1 dev set, and test sets from EMRQA and MADE 1.0. They show consistent improvements in ECE and Exact Match Accuracy.