UAT-LITE: Inference-Time Uncertainty-Aware Attention for Pretrained Transformers

TL;DR

UAT-Lite tackles the common problem of miscalibration in pretrained transformers by enabling inference-time uncertainty to shape internal computations. It estimates token-level epistemic uncertainty using Monte Carlo dropout and modulates self-attention through an uncertainty-weighted mechanism, without altering pretrained weights or training objectives. A layer-wise variance decomposition provides diagnostic insight into how predictive uncertainty propagates across depth, and a confidence-aware decision shaping strategy enables selective prediction. Empirical results across SQuAD 2.0, MNLI, SST-2, and clinical tasks show notable reductions in Expected Calibration Error (approximately 20% on average) and improved robustness under distribution shift, with modest inference-time cost compared to deep ensembles. The approach is complementary to post-hoc calibration methods and demonstrates practical gains for reliable, selective NLP deployment on encoder-based transformers.

Abstract

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.

Figures (4)

• Figure 1: Left, A: Standard inference: deterministic forward pass without uncertainty. Center, B: Output-level uncertainty estimation via stochastic inference or post-hoc calibration. Right, C:UAT-Lite: uncertainty-aware inference in which epistemic uncertainty modulates attention with diagnostic insight.
• Figure 2: Overview of UAT-Lite. Token-level epistemic uncertainty is estimated via Monte Carlo dropout at the embedding layer and used to modulate self-attention within a pretrained transformer encoder. A layer-wise variance decomposition provides diagnostic attribution of predictive uncertainty across transformer.
• Figure 3: Uncertainty-aware attention under negation (illustrative). Attention heatmaps for a standard transformer (left), UAT-Lite (middle), and their difference (right) on the sentence "The patient did not respond positively to the treatment." While absolute attention appears similar due to normalization, the difference map highlights small, structured redistributions away from sentiment-bearing tokens toward negation and predicate structure. This visualization is qualitative and shown for a representative run (single seed); it is intended to aid interpretability rather than provide a corpus-level statistic.
• Figure 4: Component ablation on SQuAD 2.0 (validation). Expected Calibration Error (ECE; lower is better) is the primary metric; accuracy is reported for reference. Uncertainty-weighted attention yields the largest ECE reduction, and the combined model achieves the strongest calibration improvement.

Theorems & Definitions (2)

• Theorem 1: Layer-Wise Variance Decomposition
• proof