TokUR: Token-Level Uncertainty Estimation for Large Language Model Reasoning

TL;DR

TokUR introduces token-level uncertainty estimation for LLM reasoning by perturbing attention-layer weights with a low-rank noise to create an approximate posterior. Token-level uncertainties (TU, AU, EU) are computed per token and aggregated to provide a response-level confidence, with theoretical guarantees that the estimator is unbiased. Empirically, TokUR signals correlate with answer correctness across GSM8K, MATH500, and DeepScaleR and improve test-time reasoning when used for candidate selection or as an implicit reward. The work extends Bayesian uncertainty estimation to long-form generation and demonstrates scalable, training-free uncertainty estimation for robust reasoning.

Abstract

While Large Language Models (LLMs) have demonstrated impressive capabilities, their output quality remains inconsistent across various application scenarios, making it difficult to identify trustworthy responses, especially in complex tasks requiring multi-step reasoning. In this paper, we propose a Token-level Uncertainty estimation framework for Reasoning (TokUR) that enables LLMs to self-assess and self-improve their responses in mathematical reasoning. Specifically, we introduce low-rank random weight perturbation during LLM decoding to generate predictive distributions for token-level uncertainty estimation, and we aggregate these uncertainty quantities to capture the semantic uncertainty of generated responses. Experiments on mathematical reasoning datasets of varying difficulty demonstrate that TokUR exhibits a strong correlation with answer correctness and model robustness, and the uncertainty signals produced by TokUR can be leveraged to enhance the model's reasoning performance at test time. These results highlight the effectiveness of TokUR as a principled and scalable approach for improving the reliability and interpretability of LLMs in challenging reasoning tasks.

Figures (10)

• Figure 1: Distribution of TokUR's Uncertainty Scores and AUROC across Different Difficulty Levels, applied to Llama-3.2-1B-Instruct. Left:TokUR (AU, Ours); Middle:TokUR (TU, Ours); Right:TokUR (EU, Ours).
• Figure 2: Performance on GSM8K (Left) and MATH500 (Right) when scaling up sample size $N$ at test time of Llama-3.2-1B-Instruct. Our TokUR (AU, EU, and TU) consistently outperforms the LL baseline, particularly when $N$ is small. Please refer to Table \ref{['tab:scaling-detailed']} for detailed numerical results.
• Figure 3: Distribution of responses from GSM8K cobbe2021gsm8k plotted in the Length Normalized EU-AU uncertainty space, as quantified by our token-level uncertainty metrics (Eqn. \ref{['eq:unc-long-estimation']}).
• Figure 4: Left: Uncertainty estimation with different perturbation strength $\sigma_q$. Right: Influence of perturbation strength on uncertainty-based AUROC scores.
• Figure 5: Left: Uncertainty estimations in different token decoding temperature $\tau$. Right: Influence of token decoding temperature on uncertainty-based AUROC scores.

Theorems & Definitions (12)

• Definition 2.1: Query-Level Uncertainty
• Definition 3.1: Response-Level Uncertainty
• Proposition 3.1: Response-Level Uncertainty as an Unbiased Estimator of Query-Level Uncertainty
• Proposition 3.2: Token-Level and Response-Level Uncertainty
• Lemma C.1: Definition of Conditional Entropy cover1999elements
• Lemma C.2: Chain rule of Conditional Entropy cover1999elements
• Proposition C.1: Decomposition of Query-Level Uncertainty, Eqn. \ref{['eq:chain-rule']}
• proof
• proof
• proof