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

• URL: http://arxiv.org/abs/2602.02952v1
• Date: Tue, 03 Feb 2026 00:51:26 GMT
• Title: UAT-LITE: Inference-Time Uncertainty-Aware Attention for Pretrained Transformers
• Authors: Elias Hossain, Shubhashis Roy Dipta, Subash Neupane, Rajib Rana, Ravid Shwartz-Ziv, Ivan Garibay, Niloofar Yousefi,
• Abstract summary: UAT-LITE is an inference-time framework that makes self-attention uncertainty-aware.<n>It reduces Expected Error by approximately 20% on average relative to a fine-tuned BERT-base baseline.<n>It improves selective prediction and robustness under distribution shift.
• Score: 11.741258610945259
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Direct Learning of Calibration-Aware Uncertainty for Neural PDE Surrogates [0.0]
  Cross-regularized uncertainty learns uncertainty parameters during training using routed through a held-out regularization split.<n>The predictor is optimized on the training split for fit, while low-dimensional uncertainty controls are optimized on the regularization split to reduce train-test mismatch.<n>We instantiate the approach in Fourier Neural Operators and evaluate on APEBench sweeps over observed fraction and training-set size.
  arXiv  Detail & Related papers  (2026-02-11T17:57:20Z)
• CATTO: Balancing Preferences and Confidence in Language Models [4.678970068275123]
  Large language models (LLMs) often make accurate next token predictions but their confidence in these predictions can be poorly calibrated.<n>We introduce a predictive calibration-aware objective that aligns predicted confidence with empirical prediction correctness.<n>We also introduce Confidence@k, a test-time scaling mechanism leveraging calibrated token probabilities for Bayes-optimal selection of output tokens.
  arXiv  Detail & Related papers  (2026-01-30T15:43:38Z)
• Uncertainty-Aware Post-Hoc Calibration: Mitigating Confidently Incorrect Predictions Beyond Calibration Metrics [6.9681910774977815]
  This paper presents a post-hoc calibration framework to enhance calibration quality and uncertainty-aware decision-making.<n>A comprehensive evaluation is conducted using calibration metrics, uncertainty-aware performance measures, and empirical conformal coverage.<n> Experiments show that the proposed method achieves lower confidently incorrect predictions, and competitive Expected Error compared with isotonic and focal-loss baselines.
  arXiv  Detail & Related papers  (2025-10-19T23:55:36Z)
• Uncertainty-Driven Reliability: Selective Prediction and Trustworthy Deployment in Modern Machine Learning [1.2183405753834562]
  This thesis investigates how uncertainty estimation can enhance the safety and trustworthiness of machine learning (ML) systems.<n>We first show that a model's training trajectory contains rich uncertainty signals that can be exploited without altering its architecture or loss.<n>We propose a lightweight, post-hoc abstention method that works across tasks, avoids the cost of deep ensembles, and achieves state-of-the-art selective prediction performance.
  arXiv  Detail & Related papers  (2025-08-11T02:33:53Z)
• Efficient Post-Hoc Uncertainty Calibration via Variance-Based Smoothing [0.810304644344495]
  We investigate whether incorporating prior information can improve uncertainty estimates in conventional deep neural networks.<n>Our focus is on machine learning tasks where meaningful predictions can be made from sub-parts of the input.
  arXiv  Detail & Related papers  (2025-03-19T16:47:15Z)
• Provably Reliable Conformal Prediction Sets in the Presence of Data Poisoning [53.42244686183879]
  Conformal prediction provides model-agnostic and distribution-free uncertainty quantification.<n>Yet, conformal prediction is not reliable under poisoning attacks where adversaries manipulate both training and calibration data.<n>We propose reliable prediction sets (RPS): the first efficient method for constructing conformal prediction sets with provable reliability guarantees under poisoning.
  arXiv  Detail & Related papers  (2024-10-13T15:37:11Z)
• Calibrated Large Language Models for Binary Question Answering [49.1574468325115]
  A well-calibrated model should produce probabilities that accurately reflect the likelihood of its predictions being correct. We propose a novel approach that utilizes the inductive Venn--Abers predictor (IVAP) to calibrate the probabilities associated with the output tokens corresponding to the binary labels.
  arXiv  Detail & Related papers  (2024-07-01T09:31:03Z)
• Uncertainty-Calibrated Test-Time Model Adaptation without Forgetting [65.21599711087538]
  Test-time adaptation (TTA) seeks to tackle potential distribution shifts between training and test data by adapting a given model w.r.t. any test sample.<n>Prior methods perform backpropagation for each test sample, resulting in unbearable optimization costs to many applications.<n>We propose an Efficient Anti-Forgetting Test-Time Adaptation (EATA) method which develops an active sample selection criterion to identify reliable and non-redundant samples.
  arXiv  Detail & Related papers  (2024-03-18T05:49:45Z)
• Distribution-free uncertainty quantification for classification under label shift [105.27463615756733]
  We focus on uncertainty quantification (UQ) for classification problems via two avenues. We first argue that label shift hurts UQ, by showing degradation in coverage and calibration. We examine these techniques theoretically in a distribution-free framework and demonstrate their excellent practical performance.
  arXiv  Detail & Related papers  (2021-03-04T20:51:03Z)
• Discriminative Jackknife: Quantifying Uncertainty in Deep Learning via Higher-Order Influence Functions [121.10450359856242]
  We develop a frequentist procedure that utilizes influence functions of a model's loss functional to construct a jackknife (or leave-one-out) estimator of predictive confidence intervals. The DJ satisfies (1) and (2), is applicable to a wide range of deep learning models, is easy to implement, and can be applied in a post-hoc fashion without interfering with model training or compromising its accuracy.
  arXiv  Detail & Related papers  (2020-06-29T13:36:52Z)
• Optimal Change-Point Detection with Training Sequences in the Large and Moderate Deviations Regimes [72.68201611113673]
  This paper investigates a novel offline change-point detection problem from an information-theoretic perspective. We assume that the knowledge of the underlying pre- and post-change distributions are not known and can only be learned from the training sequences which are available.
  arXiv  Detail & Related papers  (2020-03-13T23:39:40Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.