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

arXiv: https://arxiv.org/abs/2602.02952 · HTML: https://mlsanity.com/paper/2602.02952/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 31 references that are in the map.
- First posted on arXiv 3 Feb 2026.

## Builds on (most cited first, 22 of 22)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [Know What You Don't Know: Unanswerable Questions for SQuAD](https://mlsanity.com/paper/1806.03822/) (2018)
- [What Disease does this Patient Have? A Large-scale Open Domain Question Answering Dataset from Medical Exams](https://mlsanity.com/paper/2009.13081/) (2020)
- [PubMedQA: A Dataset for Biomedical Research Question Answering](https://mlsanity.com/paper/1909.06146/) (2019)
- [Snapshot Ensembles: Train 1, get M for free](https://mlsanity.com/paper/1704.00109/) (2017)
- [Calibration of Pre-trained Transformers](https://mlsanity.com/paper/2003.07892/) (2020)
- [Bayesian Attention Modules](https://mlsanity.com/paper/2010.10604/) (2020)
- [Bayesian Attention Belief Networks](https://mlsanity.com/paper/2106.05251/) (2021)
- [Bayesian Transformer Language Models for Speech Recognition](https://mlsanity.com/paper/2102.04754/) (2021)
- [Calibrating Transformers via Sparse Gaussian Processes](https://mlsanity.com/paper/2303.02444/) (2023)
- [Uncertainty Aware Learning for Language Model Alignment](https://mlsanity.com/paper/2406.04854/) (2024)
- [C-LoRA: Contextual Low-Rank Adaptation for Uncertainty Estimation in Large Language Models](https://mlsanity.com/paper/2505.17773/) (2025)
- [Ensembling Pruned Attention Heads For Uncertainty-Aware Efficient Transformers](https://mlsanity.com/paper/2510.18358/) (2025)
- [Uncertainty-Aware Natural Language Inference with Stochastic Weight Averaging](https://mlsanity.com/paper/2304.04726/) (2023)
- [ALIEN: Aligned Entropy Head for Improving Uncertainty Estimation of LLMs](https://mlsanity.com/paper/2505.15443/) (2025)

## Built on by (most cited first, 4 of 4)

- [Right Knowledge, Wrong Answer: Characterizing Parametric Temporal Conflict in Open-Weight Language Models](https://mlsanity.com/paper/2606.20959/) (2026)
- [OracleZoom: On-Policy Self-Distillation Inspired Reference-Constrained Recursive Image Super Resolution](https://mlsanity.com/paper/2609.06490/) (2026)
- [BEFORE THE FLIP: Measuring Hidden Score Shifts In Quantized Vision Language Models Before The Answer Changes for Visual Question Answering](https://mlsanity.com/paper/2609.06922/) (2026)
- [New Evidence, Same Choice: Testing Physical Experiment Selection in Vision Language Models](https://mlsanity.com/paper/2609.11022/) (2026)

## Abstract (arXiv)

Neural NLP models are often miscalibrated and overconfident, assigning high confidence to incorrect predictions and failing to express uncertainty during internal evidence aggregation. This 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 via Monte Carlo dropout in pretrained transformer classifiers. Unlike output-level calibration (e.g., TS), UAT-LITE injects epistemic uncertainty directly into attention, enabling uncertainty-aware routing during contextualization and token-level diagnostic signals beyond global logit rescaling. 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 layer-wise variance decomposition to diagnose how predictive uncertainty accumulates across transformer depth. Across SQuAD 2.0 answerability, MNLI, and SST-2, UAT-LITE achieves an average relative ECE reduction of approximately 20% compared with a fine-tuned BERT-base baseline while preserving accuracy, and yields more informative uncertainty behavior for selective prediction under distribution shift.