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

TL;DR

TokUR introduces a token-level uncertainty estimation method for LLMs that improves reasoning reliability and interpretability by self-assessment and self-improvement during mathematical tasks.

Contribution

The paper presents a novel low-rank random weight perturbation technique for token-level uncertainty estimation in LLMs, enhancing their reasoning performance.

Findings

TokUR's uncertainty correlates strongly with answer correctness.

Uncertainty signals improve model robustness in reasoning tasks.

Method is scalable and effective across various difficulty levels.

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.