Source-linked AI summary
Calibrating Structured Output Predictors for Natural Language Processing
Abhyuday Jagannatha, Hong Yu
TL;DR
Structured NLP outputs can be too large for standard binary or multi-class calibration, despite the importance of reliable confidence for safety-critical use. The paper calibrates confidence for selected output entities and uses those scores to re-score predictions. It reports improved calibration across tasks, improved model accuracy from re-scoring, and out-of-domain gains, while noting remaining domain-shift limitations.
Problem
Structured NLP output spaces can be too large for standard binary or multi-class calibration, although calibrated confidence is important for identifying potentially incorrect predictions in safety-critical applications.
Source: §1 · p. 1Method
The method constructs event-of-interest sets for output entities and calibrates their confidence, then uses calibrated scores to re-score predictions.
Source: §2.3 · p. 2, §1 · p. 2Results
Takeaways & Limitations
Entity-specific calibrated scores can also serve as an uncertainty-aware decoding step that improves model performance without additional training or data requirements.
Source: §1 · p. 2, §5 · p. 9The out-of-domain QA setup shifts the answer-passage distribution but keeps the questions the same, so analysis with both data and question shifts remains future work.
Source: §4 · p. 8
Abstract
from arXiv · showhide
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 system is 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.
1 Introduction
Structured NLP outputs are too large for standard binary or multi-class calibration, yet reliable confidence matters in safety-critical applications. The paper calibrates confidence for specific output entities and uses it to re-score predictions, reporting improved calibration and model performance.
1 Introduction
- Calibrated confidence can help NLP systems identify predictions likely to be incorrect, a property described as necessary for responsible use in safety-critical domains.The introduction names healthcare and finance as examples.Source: §1 · p. 1
- Structured NLP output spaces can be too large for binary or multi-class calibration methods to apply directly.The paper cites information extraction and extractive question answering as examples.Source: §1 · p. 1
- The paper studies confidence calibration for neural models used in POS, NER, and QA, including models based on contextual embeddings.Such models are being adopted for information extraction and other NLP tasks in safety-critical applications.Source: §1 · pp. 1–2
- The proposed method improves calibration across benchmark NLP tasks and out-of-domain settings, while re-scoring improves model performance across several tasks.The reported improvements require no additional training or data for the re-scoring step.Source: §1 · p. 2
- The method calibrates confidence for output entities of interest rather than requiring confidence calibration for the complete structured prediction.It provides a general scheme for constructing manageable, relevant output spaces.Source: §1 · p. 2
2 Calibration framework for Structured Prediction NLP models
The framework calibrates probabilities for selected events of interest rather than every structured output, then trains a forecaster on examples derived from those events. Its event-selection schemes balance coverage of positive events against a manageable output set, and its forecaster uses model, uncertainty, rank, and task-specific features.
2.1 Background
- Structured prediction models assign probabilities to outputs such as label sequences, but calibrating every output is difficult when the output space is large.Binary and multi-class calibration methods do not extend straightforwardly to such large structured output spaces.Source: §2.1 · p. 2
2.3 Calibration in Structured Prediction
- The framework focuses on an event-of-interest set I(x), whose events are defined from output entities relevant to a model’s deployment requirements.A forecaster maps an event, input, and model to a confidence score between 0 and 1.Source: §2.3 · pp. 2–3
2.4 Construction of Event of Interest set I(x)
- The construction of I(x) seeks a relatively small set of entities while covering as many positive events as possible.Using top-k predictions can increase positive-event coverage and training examples, but larger k also involves a trade-off.Source: §2.4 · p. 3
2.5 Forecaster Construction
- Forecaster training data is formed from entities and event-membership labels collected using top-k predictions, then used to train a gradient-boosted decision-tree forecaster.The validation data is repurposed after model training; the choice between k=2 and k=3 is selected using top-1 training data.Source: §2.5 · p. 4
2.6 Feature Construction for Calibration
- The forecaster uses mean marginal entity probability averaged over 10 MC-dropout samples, 10th- and 90th-percentile values for model-uncertainty information, and rank k as a feature for top-k MAP entity predictions.Source: §2.6 · p. 4
- Feature choices also include entity-span length for NER and language-model perplexity for QA; OOV features are excluded from the experiments.Perplexity is examined as an indicator of distributional shift, while the paper notes that word-pieces reduce the negative effect of OOV words in contextual models.Source: §2.6 · pp. 4–5
3 Experiments and Results
Experiments evaluate calibration on POS, NER, and QA, including out-of-domain QA. The reported results show improved calibration and, with confidence-based re-ranking, improved prediction performance.
3 Experiments and Results
- Experiments use BERT and distilBERT, evaluate calibration with ECE using 20 bins, and compare against Platt scaling and a calibrated-mean baseline.The calibrated-mean baseline uses gradient-boosted decision trees; both baselines use the same length-normalized, dropout-averaged probability input.Source: §3 · p. 5
3.1 Calibration for Part-of-Speech Tagging
- For POS, the entity of interest is the complete sentence-level MAP label sequence, with events indicating whether top-k predictions are correct.The study evaluates BERT, BERT-CRF, and distilBERT on Penn Treebank.Source: §3.1 · p. 5
- POS calibration methods outperform both baselines; Rank and Var features and additional top-2 prediction sequences improve calibration, while k = 3 performs worse than k = 2.Subsequent experiments report top-k calibration using the heuristic-k value.Source: §3.1 · pp. 5–6
3.2 Calibration for Named Entities
- Re-ranked NER predictions consistently improve entity-level Micro-F score over the original model predictions.The method re-scores spans from top-3 MAP predictions and removes spans below a 0.5 confidence threshold.Source: §3.2 · p. 6, Table 4 · p. 6
3.3 Calibration for QA Models
- QA experiments treat top-k answer spans as entities of interest and include open-domain, clinical-domain, and out-of-domain evaluations.MADE 1.0 QA data is constructed using EMRQA question forms, enabling evaluation with a changed text distribution.Source: §3.3 · pp. 6–7
- QA forecasters consistently improve ECE and Exact Match Accuracy across SQuAD1.1, EMRQA, and MADE 1.0, including the out-of-domain evaluation.The out-of-domain model is trained and calibrated on EMRQA; Table 5 reports ECE and Table 6 reports Exact Match Accuracy changes.Source: §3.3 · p. 7, Table 5 · p. 7, Table 6 · p. 7
4 Discussion
The discussion finds that task-agnostic forecasters perform strongly across tasks, while rank, uncertainty, and task-specific features affect calibration and re-scoring. Out-of-domain calibration improves but remains limited by prediction shifts and validation-data availability.
Features and top-k samples
- Top-k features improve performance in almost all tasks when rank is included, but omitting rank can sharply increase calibration error for CoNLL NER and MADE 1.0 QA.Rank helps distinguish candidates with similar confidence and uncertainty; in POS, rank substantially improves all models.Source: §4 · pp. 7–8
Discussion
- Task-agnostic Rank+Var+topk forecasters consistently outperform or match other forecasting methods, and proposed methods win in most tasks while remaining competitive in others.Task-specific features can reduce calibration error further.Source: §4 · pp. 7–8
Task-specific features
- The QA domain-shift experiment holds questions constant and shifts answer-passage distributions; the language-model feature does not improve out-of-domain performance in this setup.The authors call for analysis including both data and question shifts.Source: §4 · p. 8
Choice of k is important
- The best k depends on the task: heuristic-k is 2 for POS and 3 for other tasks, and these heuristic models perform best.The results suggest the Section 2.5 heuristic may generalize to other tasks.Source: §4 · p. 8
Re-scoring
- Confidence-based re-scoring modestly improves NER and QA, while POS shows no appreciable gain or drop and ineffective cases do not degrade performance.The authors suggest POS's above-97% token accuracy may explain the lack of gain; uncertainty features are useful in both calibration and re-ranking.Source: §4 · pp. 8–9
Out-of-domain Performance
- On unseen QA data, the methods improve calibration and accuracy, but a substantial in-domain versus out-of-domain calibration gap remains.The reliability plot shows the gap; out-of-domain forecasters remain over-confident, and out-of-domain validation appears to bridge it, though such data may be unavailable.Source: Figure 1 · p. 8, §4 · p. 9
Uncertainty Estimation
- Experiments use MC-Dropout for epistemic uncertainty, but the scheme is compatible with other predictive-distribution methods and does not investigate aleatoric uncertainty.The authors leave calibration with different uncertainty types for future work.Source: §4 · p. 9
5 Conclusion
The paper presents a calibration and confidence-based re-scoring scheme for structured NLP entities that outperforms competitive baselines across several tasks. The method also supports entity-specific confidence, model-accuracy improvements without additional training or data, and out-of-domain generalization.
5 Conclusion
- The calibration and confidence-based re-scoring scheme outperforms competitive baselines on several NLP tasks.The method produces calibrated outputs for specific entities rather than entire label sequences.Source: §5 · p. 9
- The method improves trained-model accuracy at no additional training or data cost.It is compatible with modern NLP architectures such as BERT.Source: §5 · p. 9
- The calibration generalizes to out-of-domain datasets without over-fitting to in-domain data.Source: §5 · p. 9
A.1 Algorithm Details:
Forecaster construction uses top-k predicted sequences to form candidate events and averages logits across MC-Dropout samples. Its forecaster data splits use top-k entities for training and top-1 entities for validation.
A.1 Algorithm Details
- Candidate events are extracted from the top-k label sequences predicted for each output.Source: §A.1 · p. 12
- Logits from the model are averaged over 10 MC-Dropout samples before the final output layer.Source: §A.1 · p. 12
- The original task validation set trains the forecaster, with top-k predicted entities in its training split and top-1 entities in its validation split.The validation set is partitioned to construct both forecaster splits.Source: §A.1 · p. 12
A.2 Evaluation Details
The evaluation uses expected calibration error (ECE), which estimates the expected absolute difference between confidence and accuracy. It bins model outputs into 20 bins and computes ECE across them.
A.2 Evaluation Details
- ECE estimates the expected absolute difference between model confidence and accuracy across output bins.Source: §A.2 · p. 12
- The experiments calculate ECE using 20 bins.Source: §A.2 · p. 12
- For each bin, accuracy and model confidence are computed; ECE uses the bin index, total sample count, and bin-level values.The passage defines N as the number of bins, n as the total number of samples, and B_i as the ith bin.Source: §A.2 · p. 12
A.3 Implementation Details
Experiments use AllenNLP with HuggingFace Transformers, initialize models with BERT-base-cased or bio-BERT weights, and apply a common learning rate. Training uses dataset validation splits when available and early stopping based on the final evaluation metric.
A.3 Implementation Details
- Implementation uses AllenNLP’s wrapper with HuggingFace Transformers and initializes models with BERT-base-cased or bio-BERT weights.Cased models are used throughout, since bio-BERT uses cased models.Source: §A.3 · p. 12
- All experiments use a learning rate of 2e-5 and provided validation splits when available.For MADE 1.0, EMRQA, and SQuAD1.1, 10% of training data is used as validation data.Source: §A.3 · p. 12
- Training uses patience 5 for early stopping, with 20,000 steps per epoch, and monitors the final evaluation metric rather than NLL.This choice is intended to reduce miscalibration of the underlying model, since Guo et al. (2017) observe that neural nets overfit on NLL.Source: §A.3 · p. 12
A.3.1 Part-of-speech experiments
The Penn Treebank POS experiments compare BERT and BERT+CRF models with calibration baselines, and report that length normalization improves calibration for both baselines and proposed models.
A.3.1 Part-of-speech experiments
- The POS evaluation uses the Penn Treebank WSJ splits, with sections 1–18 for training, 19–21 for validation, and 22–24 for testing.Source: §A.3.1 · p. 12
- Length normalization improves calibration for both baselines and proposed models.The experiments use an MC-Dropout average of length-normalized model probabilities as input to the calibration baselines.Source: §A.3.1 · p. 12
- The experiments use BERT with a linear layer and softmax, or BERT+CRF with a CRF layer over BERT-derived unary potentials.Both model choices provide the uncalibrated output probabilities.Source: §A.3.1 · p. 12
- Platt Scaling and a Gradient Boosted Decision Trees “Calibrated Mean” model use the same MC-Dropout averaged, length-normalized probability feature.The authors describe the feature’s lower variance and length invariance as making Platt Scaling a strong baseline.Source: §A.3.1 · p. 12
A.3.2 NER Experiments
The NER experiments construct candidate events from top-ranked label sequences, fit and select forecasters using validation data, and report that probability normalization improves calibration.
A.3.2 NER Experiments
- For CoNLL, “testa” is used for validation and “testb” for testing; MADE 1.0 uses a random 10% of its training data for validation.Source: §A.3.2 · p. 12
- Candidate events are built from top-k label sequences, with relevant entities extracted and converted into events of interest.Forecasters are trained across ranks up to kmax, and the model with minimum validation ECE is selected.Source: §A.3.2 · p. 13
- Length normalization improves the calibration of both baseline and proposed models.The normalized marginal probability for a span is used as input to both model types.Source: §A.3.2 · p. 13
- The experiments use BIO tags for training and also allow spans beginning with an “I-” tag during decoding.Source: §A.3.2 · p. 13
A.3.3 QA experiments
The QA experiments use SQuAD 1.1, EMRQA, and MADE QA, including a cross-corpus clinical setup; for single-span outputs, length normalization has little effect on baseline calibration.
A.3.3 QA experiments
- The QA experiments use SQuAD 1.1, EMRQA, and MADE 1.0, with EMRQA and MADE passage lengths restricted to be similar to SQuAD 1.1.Source: §A.3.3 · p. 13
- EMRQA experiments sample three question templates per annotation to limit overrepresentation of question types with many templates.Question types with fewer than three templates are up-sampled.Source: §A.3.3 · p. 13
- EMRQA and MADE QA share question forms but use clinical text from different sources, modeling QA on a new corpus.The setup is intended to represent answering fixed questions across clinical records from different sources.Source: §A.3.3 · pp. 13–14
- MADE QA is constructed by replicating EMRQA’s automated question-and-answer procedure with templates based on annotations common to both datasets.Source: §A.3.3 · p. 14
- Single-span QA always predicts start and end positions, so length normalization does not significantly affect baseline calibration; it may be more useful for non-contiguous spans with more outputs.The QA forecaster features therefore omit the “ln” feature.Source: §A.3.3 · p. 14