让预训练模型在推理时感知不确定性,提升可信度判断。
UAT-LITE: Inference-Time Uncertainty-Aware Attention for Pretrained Transformers
- 用蒙特卡洛丢弃在注意力层注入不确定性信号
- 比基线平均降低20%的预测误差校准度(ECE)
- 适合高风险场景下需要可信预测的用户
神经网络在自然语言处理中常出现校准不足和过度自信问题,对错误预测也赋予高置信度,且在内部证据聚合过程中无法表达不确定性,影响选择性预测和高风险应用。事后校准方法仅调整输出概率,未改变内部计算;集成与贝叶斯方法虽能改善不确定性,但需高昂训练或存储成本。本文提出UAT-LITE,一种推理时框架,通过在预训练变换器分类器中引入蒙特卡洛丢弃,使自注意力机制具备认知不确定性感知能力。该方法将不确定性直接注入注意力过程,实现上下文构建中的不确定性导向路由,以及超越全局逻辑重缩放的词元级诊断信号。词元级认知不确定性通过随机前向传播估计,并在上下文化过程中调节自注意力,无需修改预训练权重或训练目标。此外,我们引入分层方差分解,分析预测不确定性在变压器深度上的累积过程。在SQuAD 2.0可回答性、MNLI和SST-2数据集上,相比微调后的BERT-base基线,UAT-LITE平均降低约20%的期望校准误差(ECE),同时保持准确率,并在分布偏移下提供更有效的不确定性行为用于选择性预测。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。