UAT-LITE: Inference-Time Uncertainty-Aware Attention for Pretrained Transformers
This addresses the issue of unreliable confidence estimates in NLP models for high-stakes deployment, offering an incremental improvement over existing calibration methods.
The paper tackled the problem of miscalibrated confidence in neural NLP models by proposing UAT-LITE, an inference-time framework that uses Monte Carlo dropout to make self-attention uncertainty-aware, resulting in a 20% average reduction in Expected Calibration Error across benchmarks like SQuAD 2.0, MNLI, and SST-2 while maintaining task accuracy.
Neural NLP models are often miscalibrated, assigning high confidence to incorrect predictions, which undermines selective prediction and high-stakes deployment. Post-hoc calibration methods adjust output probabilities but leave internal computation unchanged, while ensemble and Bayesian approaches improve uncertainty at substantial training or storage cost. We propose UAT-LITE, an inference-time framework that makes self-attention uncertainty-aware using approximate Bayesian inference via Monte Carlo dropout in pretrained transformer classifiers. Token-level epistemic uncertainty is estimated from stochastic forward passes and used to modulate self-attention during contextualization, without modifying pretrained weights or training objectives. We additionally introduce a layerwise variance decomposition to diagnose how predictive uncertainty accumulates across transformer depth. Across the SQuAD 2.0 answerability, MNLI, and SST-2, UAT-LITE reduces Expected Calibration Error by approximately 20% on average relative to a fine-tuned BERT-base baseline while preserving task accuracy, and improves selective prediction and robustness under distribution shift.