Language Model Uncertainty Quantification with Attention Chain

Yinghao Li, Rushi Qiang, Lama Moukheiber, Chao Zhang

Keywords: Uncertainty Estimation, Large Language Model, Attention

TL;DR: Investigates large language model uncertainty estimation with critical reasoning token backtracking using attention chain.

Abstract:

Accurately quantifying a large language model's (LLM) predictive uncertainty is crucial for judging the reliability of its answers.
While most existing research focuses on short, directly answerable questions with closed-form outputs (e.g., multiple-choice), involving intermediate reasoning steps in LLM responses is increasingly important.
This added complexity complicates uncertainty quantification (UQ) because the probabilities assigned to answer tokens are conditioned on a vast space of preceding reasoning tokens.
Direct marginalization is infeasible, and the dependency inflates probability estimates, causing overconfidence in UQ.
To address this, we propose UQAC, an efficient method that narrows the reasoning space to a tractable size for marginalization.
UQAC iteratively constructs an "attention chain" of tokens deemed "semantically crucial to the final answer via a backtracking procedure.
Starting from the answer tokens, it uses attention weights to identify the most influential predecessors, then iterates this process until reaching the input tokens.
The resulting chain is further refined with similarity filtering and probability thresholding, which reduce the reasoning space, facilitating the approximation of the marginal answer token probabilities.
We validate UQAC on multiple reasoning benchmarks with advanced open-source LLMs, demonstrating that it consistently delivers reliable UQ estimates with high computational efficiency.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the COLM Code of Ethics on https://colmweb.org/CoE.html

Author Guide: I certify that this submission complies with the submission instructions as described on https://colmweb.org/AuthorGuide.html

Submission Number: 30