Enhancing Reasoning through Process Supervision with Monte Carlo Tree Search

TL;DR

The paper tackles the challenge of reasoning in large language models by introducing a generate-and-train framework that uses Monte Carlo Tree Search to sample step-by-step reasoning and assign per-step relative-correctness scores. These scores guide a weighted log-likelihood training objective with a KL penalty, and the data generation and fine-tuning are performed iteratively. Experiments on two mathematical reasoning benchmarks, GSM8K and MATH, show that the approach outperforms strong baselines and transfers improvements across datasets, demonstrating the value of process supervision without relying on reward models. However, the gains plateau after a few iterations, and data efficiency along with LoRA-based fine-tuning may limit ultimate improvements, suggesting avenues for deeper investigation and alternative training setups.

Abstract

Large language models (LLMs) have demonstrated their remarkable capacity across a variety of tasks. However, reasoning remains a challenge for LLMs. To improve LLMs' reasoning ability, process supervision has proven to be better than outcome supervision. In this work, we study using Monte Carlo Tree Search (MCTS) to generate process supervision data with LLMs themselves for training them. We sample reasoning steps with an LLM and assign each step a score that captures its "relative correctness," and the LLM is then trained by minimizing weighted log-likelihood of generating the reasoning steps. This generate-then-train process is repeated iteratively until convergence.Our experimental results demonstrate that the proposed methods considerably improve the performance of LLMs on two mathematical reasoning datasets. Furthermore, models trained on one dataset also exhibit improved performance on the other, showing the transferability of the enhanced reasoning ability.

Figures (1)

• Figure 1: An overview of the proposed methods.