[Paper Review] Improve Mathematical Reasoning in Language Models by Automated Process Supervision

Liangchen Luo, Yinxiao Liu+10 more|arXiv (Cornell University)|Jun 5, 2024

Business Process Modeling and Analysis6 citations

TL;DR

The paper introduces OmegaPRM, a divide-and-conquer Monte Carlo Tree Search method to automatically generate per-step process supervision data for training Process Reward Models, improving math reasoning in LLMs to 69.4% on MATH.

ABSTRACT

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 extit{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 ...

Motivation & Objective

• Motivate the need for intermediate-step feedback to improve multi-step mathematical reasoning in LLMs.
• Propose an automated, scalable data collection method for per-step supervision (Process Reward Models).
• Demonstrate that automated process supervision combined with weighted self-consistency improves math problem solving on MATH.
• Reduce reliance on human annotation and brute-force sampling for process supervision data.

Proposed method

• Introduce Process Supervision and contrast with Outcome Reward Models (ORMs).
• Propose OmegaPRM, a divide-and-conquer Monte Carlo Tree Search to efficiently locate the first error in Chain-of-Thought solutions.
• Use binary search over rollouts to identify the first incorrect step with Monte Carlo estimates, achieving O(k log M) complexity.
• Construct a state-action tree with nodes encoding question, partial solution, and rollouts; train PRMs with pointwise soft labels or alternative losses.
• Train PRMs on over 1.5 million automated process annotations to enable PRM-weighted majority voting.

Figure 1 : Example tree structure built with our proposed OmegaPRM algorithm. Each node in the tree indicates a state of partial chain-of-thought solution, with information including accuracy of rollouts and other statistics. Each edge indicates an action, i.e. , a reasoning step, from the last stat

Experimental results

Research questions

• RQ1 Can automated process supervision data collection match or exceed human-annotated process data for training PRMs?
• RQ2 Does OmegaPRM efficiently identify the first error in a Chain-of-Thought to provide high-quality PRM training signals?
• RQ3 Does integrating PRMs with weighted self-consistency improve math reasoning accuracy on standard benchmarks like MATH?
• RQ4 What is the impact of different PRM training objectives (soft vs hard labels, pairwise) on accuracy?
• RQ5 How does automated data quality and step splitting influence PRM effectiveness and reasoning performance?

Key findings

• OmegaPRM enables the collection of over 1.5 million process supervision annotations at zero human cost.
• PRMs trained on OmegaPRM data achieve 69.4% accuracy on the Hendrycks MATH benchmark, a relative 36% improvement over the base model.
• PRM-weighted majority voting with weighted self-consistency outperforms baselines using human-annotated or other automatic datasets.
• Soft-label pointwise PRM training yielded 70.1% accuracy on a small PRM test set, outperforming hard-label and pairwise options.
• Flexible, semantically agnostic step splitting produced similar step-length distributions to rule-based methods and maintained performance.

(a) Monte Carlo estimation of a prefix solution.