Source-linked AI summary
Let's reward step by step: Step-Level reward model as the Navigators for Reasoning
Qianli Ma, Haotian Zhou, Tingkai Liu, Jianbo Yuan, Pengfei Liu, Yang You, Hongxia Yang
TL;DR
Multi-step LLM reasoning can compound errors, and direct search can add substantial overhead. This paper uses PRM step feedback to guide heuristic greedy search and reports improvements over CoT in mathematics and code generation.
Problem
Multi-step LLM reasoning can compound errors, while direct BFS or DFS search can introduce redundancy and substantial reasoning overhead.
Source: §1 · pp. 1–2Method
HGS-PRM evaluates generated reasoning steps to guide heuristic path search, while mutation testing helps generate step-level reward data for code.
Source: §1 · p. 2, §2.3 · p. 5Results
HGS-PRM improves over CoT in mathematics and code generation; WizardMath-13B reaches 65.4% accuracy on GSM8K versus CoT’s 63.2%.
Source: §1 · p. 2Takeaways & Limitations
The study explores PRM feedback during decoding and releases a code-specific PRM dataset for further research.
Source: §5 · p. 9PRM precision is below 90%, and when its accuracy is insufficient, HGS-PRM might reduce result accuracy.
Source: §3.4 · p. 8
Abstract
from arXiv · showhide
Recent years have seen considerable advancements in multi-step reasoning with Large Language Models (LLMs). The previous studies have elucidated the merits of integrating feedback or search mechanisms during model inference to improve the reasoning accuracy. The Process-Supervised Reward Model (PRM), typically furnishes LLMs with step-by-step feedback during the training phase, akin to Proximal Policy Optimization (PPO) or reject sampling. Our objective is to examine the efficacy of PRM in the inference phase to help discern the optimal solution paths for multi-step tasks such as mathematical reasoning and code generation. To this end, we propose a heuristic greedy search algorithm that employs the step-level feedback from PRM to optimize the reasoning pathways explored by LLMs. This tailored PRM demonstrated enhanced results compared to the Chain of Thought (CoT) on mathematical benchmarks like GSM8K and MATH. Additionally, to explore the versatility of our approach, we develop a novel method to automatically generate step-level reward dataset for coding tasks and observed similar improved performance in the code generation tasks. Thus highlighting the robust nature of our reward-model-based approach to inference for reasoning tasks.
1 INTRODUCTION
The introduction frames multi-step reasoning as vulnerable to cascading errors and proposes PRM-guided heuristic greedy search during decoding. It reports gains over CoT on math benchmarks and extends the approach to code generation with automatically generated reward data.
1 INTRODUCTION
- The authors investigate PRM feedback at inference because multi-step LLM reasoning can compound errors and PRM has traditionally been used during training.Source: §1 · pp. 1–2
- Prior work combines feedback or search with reasoning, but reflection can loop or exceed context limits, while BFS and DFS add substantial overhead.Source: §1 · pp. 1–2
- 65.4% vs CoT’s 63.2% on GSM8K and 13.7% vs 10.4% on MATH; Code-LLaMA-Python-13B rises from 41.5% to 44.5% on HumanEval.Source: §1 · p. 2
- HGS-PRM evaluates each newly generated reasoning step, accepts or replaces it, and backtracks when a valid step cannot be found.Source: §1 · p. 2
- For code generation, the authors automatically generate step-level reward data using ASTs and mutation testing, then train a code-specific PRM.Source: §1 · p. 2
2 METHOD
The method trains a step-level PRM and uses its feedback to guide heuristic greedy search over reasoning paths. It also constructs code-step labels from ground-truth programs, AST mutations, and unit tests, while noting a limitation in the mutation scope.
2.1 PROCESS-SUPERVISED REWARD MODEL
- The reward model is trained on LLaMA-7B-SFT with PRM800k to score solution steps as positive, negative, or neutral.LLaMA-7B-SFT is first instruction-tuned on MATH question-solution pairs.Source: §2.1 · p. 3
2.2 HEURISTIC GREEDY SEARCH WITH PRM
- HGS-PRM evaluates candidate reasoning steps, advances on positive scores, uses neutral children when available, and backtracks when all children are negative.Maximum bandwidth B and iteration count T bound exploration; the authors report BFS and MCTS as more redundant and less effective, favoring greedy search for its exploration-efficiency trade-off.Source: §2.2 · p. 4, Figure 2 · p. 4
2.3 PROCESS-SUPERVISED DATA FOR CODE
- PRM-Code labels ground-truth code lines as positive and AST-mutated variants as neutral or negative according to unit-test outcomes.A mutated variant that passes unit tests receives a neutral label; one that fails receives a negative label. The authors note that focusing on atomic operators limits dataset generality.Source: §2.3 · p. 5, Figure 4 · p. 6
3 EXPERIMENT
The experiments evaluate HGS-PRM on mathematical and code-generation tasks, examine how PRM training choices affect step scoring, and analyze filtered samples and example paths.
3.1 PROCESS-SUPERVISED REWARD MODEL TRAINING
- Math-data instruction tuning improved PRM performance, and positive labels had the highest precision and recall, supporting greedy acceptance of positive steps.Neutral labels were difficult to distinguish and were treated as positive or negative; the search adjusts how neutral steps are handled to avoid misjudgments.Source: Figure 5 · p. 6, §3.1 · pp. 6–7
- The code PRM used Code-LLaMA-7B and automatically generated MBPP step-level data, and it discerned negative cases better than the math PRM.The authors suggest this stronger negative-case discrimination indicates greater potential for applying PRM in coding.Source: Figure 5 · p. 6, §3.1 · p. 7
3.2 MATH RESULT
- HGS-PRM outperformed CoT on GSM8K and MATH for general and math-specific models, with math-specific models benefiting more.MATH evaluation used 500 test problems because 4.5K test problems were in the PRM800K training set; the authors emphasize aligning language-model and reward-model capabilities.Source: Table 1 · p. 7, §3.2 · p. 7
3.3 CODE RESULT
- On HumanEval, Code-LLaMA-Python-7B improved pass@1 by 4.9% (36.6% VS 41.5%) and Code-LLaMA-Python-13B by 3.0% (41.5% VS 44.5%) over CoT.These results also surpassed the Code-LLaMA paper’s greedy decoding results by 3.1% (38.4% VS 41.5%) and 1.2% (43.3% VS 44.5%), respectively.Source: Table 2 · p. 8, §3.3 · p. 8
3.4 FURTHER ANALYSIS
- The reliability analysis warns that PRM precision below 90% could reduce HGS-PRM accuracy and proposes scoring sampled paths as an alternative.The authors report that this repeated-sampling alternative has higher computational cost than search for complex reasoning problems.Source: §3.4 · p. 8
- After PRM filtering, accuracy was 14.4% and 36.66%, versus original accuracy of 4.25% and 30.13%; repeated sampling costs more than HGS-PRM on complex problems.The analysis interprets the filtered-accuracy comparison as evidence of a correlation between PRM step feedback and answer correctness, and motivates search over sampling.Source: Table 3 · p. 8, §3.4 · p. 8
3.5 CASE STUDY
- The case study presents correct and erroneous HGS-PRM reasoning paths, including parsing, computational, and semantic errors.Source: §3.5 · p. 8
4 RELATED WORKS
Prior work developed prompting, feedback, and search strategies for LLM reasoning, while PRM research remained comparatively limited, especially for coding tasks.
4 RELATED WORKS
PRM & ORM
- PRM research includes PRM800K and PPO-based RLEIF, but prior work had focused on math, leaving coding comparatively underexplored.Source: PRM & ORM · p. 9
5 CONCLUSIONS AND FUTURE WORK
The paper applies PRM feedback during heuristic greedy search for multi-step inference and reports effectiveness on MATH and GSM8K. It also introduces a code-specific reward-data contribution using mutation testing.
5 CONCLUSIONS AND FUTURE WORK
- HGS-PRM integrates PRM feedback into heuristic greedy search for multi-step inference, with experiments on MATH and GSM8K.Source: §5 · p. 9
- The study explores PRM during decoding rather than using it only in RLHF or reject sampling.Source: §5 · p. 9
- Mutation testing supports rapid scaling of code step-level reward data, and the authors release a coding-focused PRM dataset.Source: §5 · p. 9
A HEURISTIC GREEDY SEARCH WITH PRM(HGS-PRM)
HGS-PRM searches a tree of reasoning prefixes, scoring generated steps with PRM and using forward selection or backtracking to build a complete path.
A HEURISTIC GREEDY SEARCH WITH PRM(HGS-PRM)
- Each search node represents a state containing the question and reasoning prefix, along with its current reasoning step.Source: A HEURISTIC GREEDY SEARCH WITH… · p. 15
- Nodes also store the PRM step reward, cumulative reward value, and links to their parent and children.Source: A HEURISTIC GREEDY SEARCH WITH… · p. 15
- Expansion samples a possible next step conditioned on the current state, with diversity affected by temperature, top-p, and top-k.Source: A HEURISTIC GREEDY SEARCH WITH… · p. 15
- PRM assigns each new step a reward in [−1, 0, 1], and a positive reward makes the search proceed to that child.Source: A HEURISTIC GREEDY SEARCH WITH… · p. 15
- When expansion fails or reaches the bandwidth limit, the search backs up; it selects a zero-reward child when available, otherwise its parent.Source: A HEURISTIC GREEDY SEARCH WITH… · p. 15
- The algorithm stops on an end token or after its iteration limit, then returns the current state or generates a final result from it.Source: A HEURISTIC GREEDY SEARCH WITH… · p. 15
B.1 MATH CASE STUDY
The supplied table caption identifies a MATH case study with WizardMath-13B.
B.1 MATH CASE STUDY
- The table is scoped to a MATH case study using WizardMath-13B.Source: Table 4 · p. 16
B.2 GSM8K CASE STUDY
The GSM8K case study works through a sprint-distance problem, calculating the distance per session and then the weekly total of 540 meters.
B.2 GSM8K CASE STUDY
- James runs 3 sprints of 60 meters each, totaling 180 meters per session.The case calculates this by multiplying the number of sprints by the distance of each sprint.Source: §B.2 · p. 17
- With 3 sessions per week, James runs 540 meters in a week.The weekly total is calculated by multiplying 180 meters per session by 3 sessions.Source: §B.2 · p. 17
- The case states the final answer as 540 meters per week.Source: §B.2 · p. 17
- Table 5 is identified as a GSM8K case with WizardMath-13B.Source: Table 5 · p. 17
B.3 HUMANEVAL CASE STUDY
The HumanEval case is identified as using Code-LLaMA-Python-13B.
B.3 HUMANEVAL CASE STUDY
- Table 6 is identified as a HumanEval case with Code-LLaMA-Python-13B.Source: Table 6 · p. 18
Related papers
This paper cites
- Language Models are Few-Shot Learners2020
- Training language models to follow instructions with human feedback2022
- LLaMA: Open and Efficient Foundation Language Models2023
- Self-Consistency Improves Chain of Thought Reasoning in Language Models2022
- Let's Verify Step by Step2023
- Tree of Thoughts: Deliberate Problem Solving with Large Language Models2023
Cited by
- Unlocking Token Rewards via Training-Free Reward Attribution2026
- LeanPRM: Lean-Grounded Process Supervision for Natural Language Proof Verification2026
- When LLM Meets Tree Search: A Systematic View of Inference as Search in Large Language Models2026
- Performance Foundations of Parallel & Distributed Reasoning Language Models2026
- Neuro-symbolic PRM: Enhancing Scientific Reasoning via Structured Traces and Symbolic Verification2026
- Think Again or Think Longer? Selective Verification for Budget-Aware Reasoning2026