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

arXiv:2310.10080v1cs.CL

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–2
  • Method

    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. 5
  • Results

    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. 2
  • Takeaways & Limitations

    The study explores PRM feedback during decoding and releases a code-specific PRM dataset for further research.

    Source: §5 · p. 9

    PRM precision is below 90%, and when its accuracy is insufficient, HGS-PRM might reduce result accuracy.

    Source: §3.4 · p. 8

Abstract

from arXiv · show

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

  • CoT, Zero-shot-CoT, PAL, self-consistency, and Auto-CoT are among prompting approaches used to improve LLM reasoning.Source: §4 · p. 9
  • Feedback and search methods include self-evaluation, Tree of Thoughts, planning, code-error feedback, and retrieval-based proof support.Source: §4 · p. 9

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)

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
Loading 2310.10080v1…