Source-linked AI summary
Uncertainty-Aware Step-wise Verification with Generative Reward Models
Zihuiwen Ye, Luckeciano Carvalho Melo, Younesse Kaddar, Phil Blunsom, Sam Staton, Yarin Gal
TL;DR
Outcome-only rewards can overlook flawed intermediate reasoning, and PRMs that judge steps face reliability challenges including reward hacking. This paper introduces CoT Entropy to quantify uncertainty in judge-LM step-wise verification; it achieves the highest reported performance across three uncertainty metrics and supports more reliable selective verification.
Problem
Uncertainty quantification for complex, multi-step reward modeling remains little explored, although PRMs face annotation challenges and susceptibility to reward hacking.
Source: §1 · pp. 1–2Method
Results
CoT Entropy achieves the highest performance across AUROC, AUPRC, and AU-F1C, and selective verification improves reliability for reward models.
Source: §5.1 · pp. 6–7Takeaways & Limitations
Using uncertainty to select lower-uncertainty steps improves verification reliability, particularly for complex reasoning tasks.
Source: §5.1 · p. 7The judge-LM prompt could be further optimized.
Source: §6 · p. 8
Abstract
from arXiv · showhide
Complex multi-step reasoning tasks, such as solving mathematical problems, remain challenging for large language models (LLMs). While outcome supervision is commonly used, process supervision via process reward models (PRMs) provides intermediate rewards to verify step-wise correctness in solution traces. However, as proxies for human judgement, PRMs suffer from reliability issues, including susceptibility to reward hacking. In this work, we propose leveraging uncertainty quantification (UQ) to enhance the reliability of step-wise verification with generative reward models for mathematical reasoning tasks. We introduce CoT Entropy, a novel UQ method that outperforms existing approaches in quantifying a PRM's uncertainty in step-wise verification. Our results demonstrate that incorporating uncertainty estimates improves the robustness of judge-LM PRMs, leading to more reliable verification.
1 Introduction
Outcome-only rewards can endorse flawed reasoning when a correct answer is reached by spurious steps, while PRMs provide finer-grained step-wise feedback but remain costly to annotate and vulnerable to reward hacking. This paper applies uncertainty quantification to step-wise verification and introduces CoT Entropy for judge-LM PRMs.
1 Introduction
- Outcome-only rewards can assign high reward to flawed intermediate reasoning when the final answer is correct.Such false positives erode trust in the reasoning process.Source: §1 · p. 1
- PRMs assess correctness at each reasoning step, but high-quality step-wise annotations are costly or noisy and learned proxies remain susceptible to reward hacking.Annotation approaches include human annotation, Monte Carlo sampling, and LLM-as-a-judge.Source: §1 · p. 1
- The paper proposes uncertainty quantification to detect when a verifier is uncertain at a step and improve the reliability and interpretability of its feedback.It targets the relatively little-explored setting of complex, multi-step reward modeling.Source: §1 · p. 2
- CoT Entropy estimates uncertainty from the verifier’s reasoning, and uncertainty estimates improve judge-LM PRM verification across AUC, rejection-F1, and selective performance score.The paper also decomposes predictive uncertainty into knowledge and aleatoric uncertainty to examine PRM error sources.Source: §1 · p. 2
2 Problem Statement
The paper formulates step-wise verification as judging errors in a generated reasoning trace, then quantifies uncertainty in a judge-LM’s decision from its predictive distribution. The setup distinguishes the PRM’s step-level judgments from uncertainty estimates used to identify likely mistakes.
2.1 Background on LLM Reasoning and Reward Verification
- Process reward models evaluate reasoning traces step by step, offering finer-grained supervision than outcome reward models but requiring costly step-wise labels.A PRM can flag a path that reaches the correct answer through flawed reasoning.Source: §2.1 · p. 2
2.2 Step-wise Verification with Generative PRM
- A generative PRM predicts whether each reasoning step contains an error, with the probability of an error-free step serving as its step-wise reward.The solution-level reward is the product of the step-wise probabilities, and the verifier selects the most likely class for each step.Source: §2.2 · pp. 2–3
2.3 Uncertainty Quantification with Judge-LM
- The uncertainty task maps a judge-LM’s predictive distribution over step-error labels to an estimate of uncertainty in its verification decision.The label probabilities are normalized from the token probabilities assigned to the label tokens; uncertainty estimates are assessed against ground-truth step labels.Source: §2.3 · p. 3
3 Chain-of-Thought Entropy
CoT Entropy estimates uncertainty in a judge-LM’s step-wise verification by sampling reasoning paths and marginalizing their judgments. Unlike majority-vote methods, it uses the approximate posterior predictive distribution across decisions.
3 Chain-of-Thought Entropy
- CoT prompting asks the judge-LM to produce a rationale before deciding whether a step contains an error, integrating reasoning into the verification decision.The rationale is intended to help identify subtle errors and keep the assessment focused on local logical consistency.Source: §3 · pp. 3–4
- CoT Entropy samples rationales and verification decisions, groups outputs by whether the step is judged correct or incorrect, then aggregates class probabilities to estimate entropy.The method marginalizes over sampled rationales to approximate the entropy of the verification judgment.Source: §3 · p. 4
- The method uses the full approximate posterior predictive distribution over verification outcomes, unlike self-consistency’s majority vote over samples.The context x≤t comprises the question, the reasoning trace through step t, and the CoT prompt.Source: §3 · p. 4
4 Uncertainty Quantification on Math Reasoning Verification
The experimental section evaluates uncertainty estimates for step-wise verification of mathematical reasoning, using PRM800K data and a judge-LM setup.
4.1 Datasets
- The evaluation uses PRM800K labels for GPT-4-generated MATH solutions, sampling 150 of 500 test questions (1,152 steps; 129 errors, 11.2%).The subset contains an average of 7.7 steps per solution.Source: §4.1 · pp. 4–5
4.2 Models
- QWEN2-MATH-72B-INSTRUCT3 serves as judge-LM, generating one low-temperature predicted label and 10 high-temperature reasoning paths and decisions per step.The sampling temperatures are 0.1 and 1.0, respectively.Source: §4.2 · p. 5
4.3 Baseline Methods
- Baselines are Naive Entropy, P(True), SEU with MiniLM or NV-Embed, discrete CoT Entropy, and random scores, evaluated over five seeds.Performance is reported as the mean and standard deviation across seeds.Source: §4.3 · p. 5
4.4 Evaluation metrics
- AUROC and AUPRC assess uncertainty ranking, with AUPRC's uninformative baseline set by the proportion of correctly verified steps.AUROC ranges from 0 to 1, with 0.5 representing an uninformative classifier.Source: §4.4 · p. 5
- Selective evaluation uses F1; AU-F1C summarizes retained-step F1 from 30% onward, while Rejection-F1 is reported from 60% to 100%.F1 is used given the 11.2% positive-label rate; Rejection-F1 at 100% is equivalent to an uninformative classifier.Source: §4.4 · pp. 5–6
5 Evaluation Results
On a PRM800K test-set subset, CoT Entropy leads the reported uncertainty methods for step-wise verification, and selective verification improves reliability. Predictive uncertainty performs best at identifying verifier errors, with epistemic uncertainty nearly as effective.
5.1 Does CoT Entropy inform confidence in verification?
- CoT Entropy achieves the highest performance across AUROC, AUPRC, and AU-F1C.Results are averaged over five runs.Source: §5.1 · p. 6, Table 1 · p. 6
- CoT Entropy achieves the highest Rejection-F1 across thresholds, while Rejection-F1 generally rises as more examples are rejected.The figure plots retained-example F1 against rejection threshold.Source: Figure 1 · p. 6, §5.1 · p. 7
- Naive entropy without clustering performs worse than random; discrete CoT Entropy trails its probabilistic variant, unlike semantic entropy’s small variant gap.The authors hypothesize that probabilities capture reasoning-path nuances in math verification.Source: §5.1 · p. 7
- NV-EMBED outperforms MINILM among embedding-based methods, while CoT Entropy outperforms SEU, which relies on an external embedding model.The authors suggest judge-LM token probabilities can quantify uncertainty without external-model heuristics when the predictive distribution is meaningful.Source: §5.1 · p. 7
- Selective verification of lower-uncertainty steps improves reward-model verification reliability, although Rejection-F1 leaves room for improvement.Reward verification on reasoning remains challenging.Source: §5.1 · p. 7
5.2 How does different uncertainty types perform in the verification task?
- The analysis decomposes each data point’s total predictive uncertainty into epistemic and aleatoric components.Equation 9 defines epistemic uncertainty as mutual information between step-wise verification and rationales; the rationale posterior is approximated by Monte Carlo sampling.Source: §5.2 · p. 7
- All uncertainty sources outperform a random baseline, and predictive uncertainty has the best Rejection-F1 for identifying verifier mistakes.This reflects the verification goal of flagging any potential source of prediction error.Source: §5.2 · p. 7, Figure 2 · p. 8
- Epistemic uncertainty correlates with verifier errors better than aleatoric uncertainty and is almost as effective as total uncertainty.The authors associate most identified math-verification errors with model knowledge uncertainty rather than label noise.Source: §5.2 · p. 8, Figure 2 · p. 8
6 Discussion
The paper applies uncertainty quantification to step-wise verification on mathematical reasoning tasks and proposes CoT Entropy, which outperforms baseline UQ methods. It identifies prompt optimization as a limitation and proposes using uncertainty estimates in future work on reward hacking and inference-time search.
6 Discussion
- CoT Entropy informs a judge-LM’s confidence in step-wise verification and outperforms other baseline UQ methods.The work applies UQ to generative reward models on math reasoning tasks.Source: §6 · p. 8
- The prompt for the judge-LM can be further optimized.Source: §6 · p. 8
- Future work will investigate using estimated uncertainty to potentially mitigate reward hacking during RL training and aid inference-time search.The authors also plan to fine-tune a generative reward model to assess whether it produces more accurate uncertainty estimates.Source: §6 · p. 8
A.1 Prompt Template
The appendix presents a prompt template for step-wise verification and reports that prompt choice matters for JSON parsing success, which serves as an indicator of the model’s understanding of the task.
A.1 Prompt Template
- The appendix provides a template for prompting LLMs to perform step-wise verification.Source: §A.1 · p. 10
- The authors tested multiple prompting templates and found that it is crucial to use one that increases JSON parsing success.Source: §A.1 · p. 10
- JSON parsing success serves as an indicator of the model’s understanding of the verification task.Source: §A.1 · p. 10
A.2 Verification Performance
The comparison reports higher accuracy for No-CoT, largely because it predominantly predicts the majority label of no error. Since errors are the minority class, the passage identifies F1 as a more meaningful measure of error detection.
A.2 Verification Performance
- No-CoT achieves higher accuracy, largely by predominantly predicting label 0 (no error).Source: Table 2 · p. 10
- Errors constitute the minority class in the imbalanced dataset.Source: Table 2 · p. 10
- F1 score is more meaningful than accuracy here because it better reflects the model’s ability to detect errors.Source: Table 2 · p. 10
Template for Step-wise Verification
The template asks a judge to assess whether the next proposed step contains an error relative to preceding steps. It directs the judge to check calculations and logical progression, then return concise reasoning and a yes/no evaluation in JSON.
Template for Step-wise Verification
- The judge evaluates the next proposed step for errors in relation to the preceding steps.The template supplies the problem, solution so far, and next step as inputs.Source: Template for Step-wise Verification · p. 11
- The criteria require checking calculations and logical progression, while focusing only on the immediate transition.Logical progression is omitted for the first step, and future actions should not be treated as errors in the current step.Source: Template for Step-wise Verification · p. 11
- The response format requests concise reasoning followed by a yes/no error judgment in JSON.Source: Template for Step-wise Verification · p. 11