Reasoning with Confidence: Efficient Verification of LLM Reasoning Steps via Uncertainty Heads
Jingwei Ni, Ekaterina Fadeeva, Tianyi Wu, Mubashara Akhtar, Jiaheng Zhang, Elliott Ash, Markus Leippold, Timothy Baldwin, See-Kiong Ng, Artem Shelmanov, Mrinmaya Sachan
TL;DR
The paper tackles the reliability and efficiency of verifying multi-step LLM reasoning by replacing costly process reward models with lightweight, data-driven uncertainty heads that exploit internal LLM states. By training UHeads on top of frozen LLMs and obtaining automatic labels from an external verifier or self-supervision, the approach achieves competitive or superior step-level verification and test-time decoding performance across math, planning, and QA tasks while using orders of magnitude fewer parameters. The results show strong generalization to out-of-domain tasks and indicate complementary gains when combining UHeads with PRMs, suggesting a scalable path toward introspective, scalable reasoning systems for LLMs. The work emphasizes efficiency, broad applicability, and the promise of internal uncertainty signals as reliable indicators of reasoning quality in large language models.
Abstract
Solving complex tasks usually requires LLMs to generate long multi-step reasoning chains. Previous work has shown that verifying the correctness of individual reasoning steps can further improve the performance and efficiency of LLMs on such tasks and enhance solution interpretability. However, existing verification approaches, such as Process Reward Models (PRMs), are either computationally expensive, limited to specific domains, or require large-scale human or model-generated annotations. Thus, we propose a lightweight alternative for step-level reasoning verification based on data-driven uncertainty scores. We train transformer-based uncertainty quantification heads (UHeads) that use the internal states of a frozen LLM to estimate the uncertainty of its reasoning steps during generation. The approach is fully automatic: target labels are generated either by another larger LLM (e.g., DeepSeek R1) or in a self-supervised manner by the original model itself. UHeads are both effective and lightweight, containing less than 10M parameters. Across multiple domains, including mathematics, planning, and general knowledge question answering, they match or even surpass the performance of PRMs that are up to 810x larger. Our findings suggest that the internal states of LLMs encode their uncertainty and can serve as reliable signals for reasoning verification, offering a promising direction toward scalable and generalizable introspective LLMs.
