[Paper Review] Calibration of Encoder Decoder Models for Neural Machine Translation

Aviral Kumar, Sunita Sarawagi|arXiv (Cornell University)|Mar 3, 2019

Natural Language Processing Techniques32 references40 citations

TL;DR

The paper analyzes calibration of attention-based encoder-decoder NMT models, identifies EOS and attention uncertainty as key miscalibration sources, and proposes a context-aware recalibration method that improves token- and sequence-level calibration and stabilizes BLEU across beam sizes.

ABSTRACT

We study the calibration of several state of the art neural machine translation(NMT) systems built on attention-based encoder-decoder models. For structured outputs like in NMT, calibration is important not just for reliable confidence with predictions, but also for proper functioning of beam-search inference. We show that most modern NMT models are surprisingly miscalibrated even when conditioned on the true previous tokens. Our investigation leads to two main reasons -- severe miscalibration of EOS (end of sequence marker) and suppression of attention uncertainty. We design recalibration methods based on these signals and demonstrate improved accuracy, better sequence-level calibration, and more intuitive results from beam-search.

Motivation & Objective

• Identify and quantify calibration gaps in state-of-the-art attention-based NMT systems (Bahdanau, GNMT, Transformer).
• Diagnose root causes of miscalibration, focusing on EOS token behavior and attention uncertainty.
• Develop a recalibration method that uses input coverage and attention entropy to adjust token probabilities.
• Demonstrate improvements in token-level calibration, sequence-level calibration, and beam-search stability.
• Show that calibration can yield BLEU gains without altering training procedures.

Proposed method

• Measure token-level calibration with weighted Expected Calibration Error (weighted ECE) across six NMT benchmarks.
• Analyze token-wise calibration, emphasizing EOS and attention-uncertainty signals.
• Propose a two-part recalibration: (i) an EOS correction as a function of input coverage, (ii) a neural-network based, input-context dependent temperature scaling for logits.
• Train calibration parameters on a held-out validation set and apply to pre-trained models.
• Compare against single-temperature scaling and assess impact on BLEU and beam-size robustness.
• Use a structured BLEU-based calibration measure (Structured ECE) to evaluate sequence-level calibration.

Experimental results

Research questions

• RQ1 How calibrated are token-level probabilities in contemporary encoder-decoder NMT models under teacher forcing?
• RQ2 What are the main factors driving miscalibration in NMT, and how do EOS signaling and attention uncertainty contribute?
• RQ3 Can a context-aware recalibration scheme improve both token-level and sequence-level calibration without retraining the models?
• RQ4 Does improved calibration stabilize beam-search inference and improve BLEU across varying beam widths?

Key findings

• All six evaluated NMT models are miscalibrated, with ECE values ranging from 2.9 to 9.8 and over-confidence observed in several cases.
• EOS calibration is particularly poor, often over- or under-estimated depending on model type, indicating weak EOS signaling.
• Attention uncertainty is linked to higher calibration error, especially for high-entropy attention distributions.
• Tail/head token calibration varies across models, showing that miscalibration affects the full softmax output, not just the top predictions.
• A context-aware recalibration method (coverage-informed EOS correction plus entropy/logit-based temperature scaling) reduces ECE across models, often improving BLEU by up to 0.4 points on several tasks.