Source-linked AI summary
Improve Mathematical Reasoning in Language Models by Automated Process Supervision
Liangchen Luo, Yinxiao Liu, Rosanne Liu, Samrat Phatale, Meiqi Guo, Harsh Lara, Yunxuan Li, Lei Shu, Yun Zhu, Lei Meng, Jiao Sun, Abhinav Rastogi
TL;DR
Complex multi-step reasoning exposes limits of final-outcome verification and makes process-supervision data expensive to collect. The paper introduces OmegaPRM, a divide-and-conquer MCTS method for automated process supervision, and reports improved mathematical reasoning across MATH500 and GSM8K.
Problem
Final-outcome reward models do not properly reward or penalize intermediate outcomes in lengthy reasoning chains, while existing process-supervision data collection is expensive to scale.
Method
OmegaPRM uses divide-and-conquer Monte Carlo Tree Search to automate process-supervision data collection by pinpointing the first Chain-of-Thought error and balancing positive and negative examples.
Results
OmegaPRM supervision combined with weighted self-consistency improves mathematical reasoning performance for Gemini Pro and Gemma2 27B on MATH500 and GSM8K.
Takeaways & Limitations
The method enables collection of over 1.5 million process-supervision annotations and improves reported LLM mathematical reasoning performance.
Takeaways & Limitations
The method requires question and golden-answer pairs, limiting its applicability to tasks with such pairs and requiring further adaptation for open-ended tasks.
Abstract
from arXiv · showhide
Complex multi-step reasoning tasks, such as solving mathematical problems or generating code, remain a significant hurdle for even the most advanced large language models (LLMs). Verifying LLM outputs with an Outcome Reward Model (ORM) is a standard inference-time technique aimed at enhancing the reasoning performance of LLMs. However, this still proves insufficient for reasoning tasks with a lengthy or multi-hop reasoning chain, where the intermediate outcomes are neither properly rewarded nor penalized. Process supervision addresses this limitation by assigning intermediate rewards during the reasoning process. To date, the methods used to collect process supervision data have relied on either human annotation or per-step Monte Carlo estimation, both prohibitively expensive to scale, thus hindering the broad application of this technique. In response to this challenge, we propose a novel divide-and-conquer style Monte Carlo Tree Search (MCTS) algorithm named \textit{OmegaPRM} for the efficient collection of high-quality process supervision data. This algorithm swiftly identifies the first error in the Chain of Thought (CoT) with binary search and balances the positive and negative examples, thereby ensuring both efficiency and quality. As a result, we are able to collect over 1.5 million process supervision annotations to train Process Reward Models (PRMs). This fully automated process supervision alongside the weighted self-consistency algorithm is able to enhance LLMs' math reasoning performances. We improved the success rates of the instruction-tuned Gemini Pro model from 51\% to 69.4\% on MATH500 and from 86.4\% to 93.6\% on GSM8K. Similarly, we boosted the success rates of Gemma2 27B from 42.3\% to 58.2\% on MATH500 and from 74.0\% to 92.2\% on GSM8K. The entire process operates without any human intervention or supervision, making our method both financially and ...
1. Introduction
The paper motivates process supervision for complex multi-step reasoning and introduces OmegaPRM, an automated divide-and-conquer MCTS algorithm for collecting process supervision data. It combines this supervision with weighted self-consistency to improve mathematical reasoning performance.
- Motivation: Process Reward Models provide feedback at each reasoning step, addressing a limitation of Outcome Reward Models that only evaluate final outcomes.The paper distinguishes granular process feedback from end-of-process outcome feedback for complex reasoning tasks.
- Motivation: Human annotation is costly, while per-step Monte Carlo estimation remains inefficient because of the large reasoning search space.These costs create the primary bottleneck in obtaining process supervision signals for training PRMs.
- OmegaPRM: OmegaPRM is a divide-and-conquer Monte Carlo Tree Search algorithm for automated process supervision data generation.For each question, it builds a Monte Carlo tree to support process annotation.
- OmegaPRM: OmegaPRM collects over 1.5 million process supervision annotations without human annotation.The authors describe the resulting dataset as the largest and highest quality dataset of its kind to date.
- Results: Weighted self-consistency with the resulting process supervision improves Gemini Pro success rates from 51% to 69.4% on MATH500 and from 86.4% to 93.6% on GSM8K.The reported gains use the instruction-tuned Gemini Pro model.
2. Related Work
Prior work improves mathematical reasoning with prompting, supervised fine-tuning, reward models, and Monte Carlo estimation. The paper positions OmegaPRM as a further use of MCTS for collecting process supervision data.
- Mathematical reasoning: Chain-of-thought prompting and supervised fine-tuning with question-response pairs containing full reasoning steps improve mathematical reasoning performance.These approaches target more complicated reasoning than pretrained models solve reliably.
- Reward models: Outcome Reward Models and Process Reward Models both improve mathematical reasoning over self-consistency, while evidence indicates that PRMs outperform ORMs.PRMs provide feedback on intermediate reasoning steps rather than only final outcomes.
- Automated supervision: Math-Shepherd and MiPS automate process-data collection with Monte Carlo estimation and report large performance gains with human involvement.The paper describes OmegaPRM as sharing their core motivation while using MCTS for process-data collection.
- Monte Carlo Tree Search: MCTS has achieved strong results in reinforcement learning through systems including AlphaGo and AlphaGo Zero, and has also been explored for LLM decoding.The paper builds on this broader use of MCTS in reasoning-related settings.
3. Methods
The method automates process supervision by using Monte Carlo rollouts to assess intermediate reasoning, binary search to locate first errors, and MCTS to reuse rollouts for PRM training. It then converts constructed tree edges and rollout statistics into labeled training examples.
- 3.1. Process Supervision: Outcome Reward Models assess only final-answer correctness, whereas Process Reward Models predict correctness for each intermediate reasoning step.PRMs therefore provide finer-grained feedback and can identify error locations in solutions.
- 3.2. Process Annotation with Monte Carlo Method: Monte Carlo annotation samples completions from a question and solution prefix to estimate whether the prefix can lead to a correct final answer.The completer policy generates rollout continuations from each prefix.
- 3.2. Process Annotation with Monte Carlo Method: Binary search locates the first incorrect step with O(klog M) policy calls instead of O(kM) for brute-force stepwise annotation.Each midpoint rollout narrows the candidate error region to one half of the solution.
- 3.3. Monte Carlo Tree Search: MCTS selects rollouts using tree statistics, applies binary search, and adds intermediate states and correctness estimates to the tree for subsequent searches.Rollouts with intermediate Monte Carlo estimates are retained as selection candidates rather than discarded.
- 3.3. Monte Carlo Tree Search: OmegaPRM constructs a state-action tree whose nodes contain question prefixes and rollouts, while edges represent single steps or consecutive step sequences.The language model supplies actions and concatenation defines state transitions.
- 3.3. Monte Carlo Tree Search: The rollout-selection heuristic prioritizes supposed-to-be-correct wrong-answer rollouts and penalizes excessively long rollouts.Supposed-to-be-correct states have Monte Carlo estimates near 1, while the length penalty discourages long searches.
- 3.4. PRM Training: Each single-step tree edge can become a PRM training example, using either Monte Carlo soft labels or binary hard labels for correctness.The PRM is trained with a standard classification loss, and the cited formulation also supports normalized pairwise preferences among alternative actions.
4. Experiments
Experiments evaluate OmegaPRM on MATH500 and GSM8K, examine its step-splitting strategy and PRM training objectives, and compare its data-generation efficiency with brute force. OmegaPRM produces strong PRM-weighted majority-voting performance while generating substantially more annotations under the same computational budget.
- 4.1. Main Results: OmegaPRM-weighted majority voting achieves 69.4% on MATH500 and 93.6% on GSM8K for fine-tuned Gemini Pro.For pretrained Gemma2 27B, it achieves 58.2% on MATH500 and 92.2% on GSM8K.
- 4.1. Main Results: As the number of sampled solutions increases, other PRMs converge toward majority-vote performance while OmegaPRM retains a clear accuracy margin.
- 4.2. Step Distribution: Most PRM800K and Math-Shepherd solutions contain fewer than 20 steps, informing OmegaPRM's binary-search splitting design.OmegaPRM divides a full solution into 16 pieces and stops when a step is shorter than the expected length.
- 4.2. Step Distribution: OmegaPRM's flexible step-splitting strategy yields a step-length distribution similar to rule-based strategies.It treats any consecutive token sequence as a valid step rather than relying only on newline delimiters.
- 4.3. PRM Training Objectives: Pointwise soft-label training gives the best per-step correctness classification accuracy among the evaluated objectives, reaching 70.1%.The comparison includes pointwise hard-label and pairwise losses on a process-supervision test set built from the MATH test split.
- 4.4. Algorithm Efficiency: With the same computational budget, OmegaPRM generates 15 million data points versus 200K for brute force, a 75-times efficiency improvement.The authors randomly down-sample the OmegaPRM data to 1.5 million annotations for PRM training.
5. Limitations
The paper identifies noise in automatic process annotations and continued dependence on question–golden-answer pairs as limitations. It proposes future comparisons and hybrid supervision, while noting the resulting task-scope constraint.
- Automatic process annotations introduce false positives and negatives, although the resulting PRM still outperforms one trained on PRM800K.The precise impact of annotation noise remains uncertain.
- Future work should compare human and automated annotations and investigate combining them for more robust and efficient supervision.
- The method requires question and golden-answer pairs to start MCTS and determine rollout correctness.This requirement limits application to tasks with such paired data.
6. Conclusion
OmegaPRM automates process-supervision data collection by locating the first CoT error with divide-and-conquer MCTS and balancing data quality. With weighted self-consistency, it improves mathematical reasoning while reducing collection costs.
- OmegaPRM uses divide-and-conquer MCTS to automate process-supervision data collection and pinpoint the first CoT error.
- Over 1.5 million process-supervision annotations are collected and used to train a PRM.
- 69.4% success on MATH, an 18.4% absolute increase over the base model, is achieved with automated supervision and weighted self-consistency.The reported relative improvement is 36%.
- The method significantly reduces data-collection costs compared with human annotation and brute-force Monte Carlo sampling.
A. Question Filtering
The paper notes that Monte Carlo evaluation of partial-solution correctness can introduce false negatives or false positives depending on question difficulty. Filtering questions is proposed to reduce this noise.
- Partial-solution MC estimation can produce false negatives when hard questions yield no correct rollout despite a correct partial solution.
- It can produce false positives when easy questions reach a correct answer despite a wrong step.
- The method cannot eliminate this noise completely, but filtering questions can reduce its occurrence.
B. Pairwise Loss Formula
The pairwise loss converts Monte Carlo values into preference probabilities for two actions, assigning ties equal probability before training the reward model.
- Bradley-Terry pairwise training requires two probability scalars that sum to 1.
- The method handles pairs with one zero and one positive MC value separately from pairs where both MC values are positive.The latter case requires normalization.
- MC values p and q are modeled as Bernoulli probabilities for actions X and Y.
- Equal outcomes are treated as ties, with each action receiving half the preference probability.
- The resulting preference probabilities are 1/2 ∗ (1 + p − q) for X and 1/2 ∗ (1 + q − p) for Y.
- After normalization, the preference probabilities can be used in the pairwise loss.