PaperlayerView exact revision on arXiv

Source-linked AI summary

ReST-MCTS*: LLM Self-Training via Process Reward Guided Tree Search

Dan Zhang, Sining Zhoubian, Ziniu Hu, Yisong Yue, Yuxiao Dong, Jie Tang

arXiv:2406.03816v3cs.CL

TL;DR

Outcome-based filtering can retain incorrect intermediate reasoning, while reliable process-reward training typically needs dense step-level annotations. ReST-MCTS* uses rollout-inferred rewards to guide MCTS* and jointly self-train policy and process reward models, reporting stronger reasoning and self-training results than prior approaches.

  • Problem

    Outcome-based self-training can include incorrect intermediate reasoning despite correct final answers, while reliable process reward models generally require dense per-step human annotations.

  • Method

    ReST-MCTS* uses rollout-inferred per-step quality values to guide MCTS* and jointly train the policy and process reward models.

  • Results

    ReST-MCTS* outperforms prior self-training approaches and achieves higher accuracy than previous reasoning baselines under the same search budget.

  • Takeaways & Limitations

    The approach supports iterative self-training with high-quality reasoning traces and inferred process-reward targets.

  • Takeaways & Limitations

    The paper has not established generalization beyond mathematical reasoning or to tasks without ground-truth answers, and calls for scaling the value model and improving filtering.

Abstract

from arXiv · show

Recent methodologies in LLM self-training mostly rely on LLM generating responses and filtering those with correct output answers as training data. This approach often yields a low-quality fine-tuning training set (e.g., incorrect plans or intermediate reasoning). In this paper, we develop a reinforced self-training approach, called ReST-MCTS*, based on integrating process reward guidance with tree search MCTS* for collecting higher-quality reasoning traces as well as per-step value to train policy and reward models. ReST-MCTS* circumvents the per-step manual annotation typically used to train process rewards by tree-search-based reinforcement learning: Given oracle final correct answers, ReST-MCTS* is able to infer the correct process rewards by estimating the probability this step can help lead to the correct answer. These inferred rewards serve dual purposes: they act as value targets for further refining the process reward model and also facilitate the selection of high-quality traces for policy model self-training. We first show that the tree-search policy in ReST-MCTS* achieves higher accuracy compared with prior LLM reasoning baselines such as Best-of-N and Tree-of-Thought, within the same search budget. We then show that by using traces searched by this tree-search policy as training data, we can continuously enhance the three language models for multiple iterations, and outperform other self-training algorithms such as ReST$^\text{EM}$ and Self-Rewarding LM. We release all code at https://github.com/THUDM/ReST-MCTS.

1 Introduction

Existing self-training can select traces with correct final answers despite incorrect intermediate reasoning, while dense step-level reward annotation is costly. ReST-MCTS* addresses this with rollout-inferred process rewards and reward-guided search, reporting stronger reasoning and self-training performance than prior approaches.

  • Motivation: Correct final answers can conceal wrong or useless intermediate reasoning, producing false-positive training traces that limit fine-tuning on complex reasoning tasks.The paper identifies inaccurate plans and intermediate steps as a central weakness of outcome-based filtering.
  • Motivation: Dense human annotations for every reasoning step are generally required to train reliable process reward models, limiting scalability.This motivates automatically deriving per-step supervision from search rather than manual labeling.
  • Approach: ReST-MCTS* uses modified MCTS guided by a trained per-step process reward model and automatically generates per-step labels through sufficient rollouts.The inferred labels both support reward-model training and help filter high-quality reasoning traces without additional human intervention.
  • Results: ReST-MCTS* outperforms prior self-training approaches, including ReSTEM and Self-Rewarding, and prior reasoning policies such as CoT and ToT.The paper reports these comparisons across multiple reasoning benchmarks.
  • Results: The reward generator produces a higher-quality process reward model than previous process reward generation techniques such as MATH-SHEPHERD.This result is reported in Table 3.
  • Results: Under the same search budget, MCTS* achieves higher accuracy than Self-Consistency and Best-of-N.The comparison is reported in Figure 2.

2 Background on Reasoning & Self-Training

LLM reasoning self-training alternates between generating traces and using a learning signal to improve the policy. Reliable tree search and self-training require value estimates, but step-level process reward supervision is difficult to obtain.

  • Reasoning setup: A reasoning policy π generates autoregressive traces of reasoning steps for an input problem Q, with the final step treated as the answer.The setup models each reasoning step as a single sentence containing multiple tokens.
  • Reasoning policies: Self-Consistency samples multiple traces and selects the most frequent final answer, whereas Best-of-N selects the trace with the highest learned value.The learned value may come from a process or outcome reward model.
  • Tree search: Tree-structured reasoning branches from intermediate steps and requires a value function to guide the otherwise combinatorially large search.Outcome reward models judge final-answer correctness, while process reward models evaluate individual reasoning steps.
  • Self-training: Self-training consists of generation, where multiple traces are sampled, followed by improvement, where a learning signal is used to fine-tune the policy; the process can iterate.The paper considers tree-structured traces in the generation stage.
  • Limitation of prior works: Reliable self-training ideally needs dense feedback on every intermediate step, but process reward models historically lack supervised per-step annotations.Sparse signals create a credit-assignment challenge similar to reinforcement learning.

3 The ReST-MCTS∗Method

ReST-MCTS* combines rollout-derived quality values, process-reward-guided MCTS*, and iterative policy and reward-model self-training. Its quality estimates reflect both correctness and progress toward a correct answer while avoiding reliance on separate outcome-reward evaluation during search.

  • Method components: ReST-MCTS* comprises MCTS*, a process reward model, a policy model, and self-training that updates both models from searched traces.MCTS* guides search with the process reward model; the policy generates intermediate reasoning steps.
  • Quality estimation: Higher v_k values indicate partial solutions more likely to be complete and correct, while v_k also captures correctness and incremental progress from individual steps.The formulation is designed to reward correct, contributing steps rather than merely correct final outputs.
  • Quality estimation: Reasoning distance m_k is the minimum number of policy reasoning steps needed to reach the correct answer from a partial solution, capturing progress and difficulty.Because m_k cannot be directly calculated, the method estimates it through simulations or trace sampling.
  • Quality estimation: The weighted reward w_s^k combines the process reward r_s^k with reasoning distance to measure a step’s incremental contribution.The quality value v_k is updated iteratively from the previous quality value and the current weighted reward.
  • Quality estimation: The quality value satisfies v_k ∈ [0, 1], and it reaches its upper bound only when the reasoning reaches a correct answer.The boundedness result assumes the process reward sigmoid score lies in [0, 1].
  • Search-based reasoning policy: MCTS* uses quality values both as targets for the process reward model and as guidance for tree search, replacing separate final-solution outcome evaluation in the proposed variant.The approach is motivated by the cost and inefficiency of fully generating every search trace for outcome-reward evaluation.
  • Search-based reasoning policy: MCTS* evaluates partial solutions, selects and backpropagates through intermediate nodes, and uses node selection, thought expansion, greedy Monte Carlo rollout, and value backpropagation.This avoids requiring every branch to be completely generated before evaluation.
  • Self-training pipeline: ReST-MCTS* iteratively searches for high-quality solutions and uses the resulting tree data to self-improve the policy and process reward model.The paper describes this as MuZero-style learning for LLMs.

4 Experiments

The experiments evaluate ReST-MCTS* as a self-training framework across policy search, iterative model improvement, and process-reward verification. Results indicate stronger search accuracy, continuous self-improvement, and higher-quality reward modeling than several baselines.

  • Reasoning Policy: Under the same search budget, MCTS* achieves higher accuracy than Self-Consistency and Best-of-N on MATH and SciBench.The broader tree-search evaluation also compares ReST-MCTS* with CoT and ToT on scientific-reasoning benchmarks.
  • Initialization of Value Model: ReST-MCTS* automatically generates per-step reward labels through sufficient MCTS* rollouts, avoiding additional human labeling for value-model training.The inferred labels are used to filter higher-quality samples for self-training.
  • Self-Training: ReST-MCTS* continuously improves three policy backbones across multiple self-training iterations and generally outperforms ReSTEM and Self-Rewarding.The comparison covers in-distribution and out-of-distribution benchmarks under three language-model backbones.
  • Self-Improvement: 48.5% accuracy exceeds Self-Consistency's 42.5% after further self-training on MATH under the same token usage.The enhanced value model enables verification to outperform Self-Consistency across the evaluated points after additional iterations.
  • Reward Guidance: ReST-MCTS* produces higher-accuracy verification than MATH-SHEPHERD and SC+MS on both GSM8K and MATH500.The experiment uses 256 outputs and evaluates Mistral-7B: MetaMATH.

5 Related Work

Related work spans data collection, prompting, supervised learning, reinforcement learning, and reasoning-time methods for improving LLM capabilities. CoT and ToT represent prompt-based and planning-based approaches to LLM reasoning.

  • LLM Improvement: Recent LLM improvement methods collect larger or higher-quality domain data, design elaborate prompts, or use supervised and reinforcement learning.These approaches are situated within broader efforts to improve reasoning and other language tasks.
  • Reasoning Algorithms: Chain-of-thought uses prompt-based reasoning, whereas tree-of-thought represents planning-based reasoning among LLM reasoning algorithms.Prior work also explores step-by-step solution generation to improve over direct generation.

6 Conclusion

ReST-MCTS* combines process-reward-guided tree search with iterative self-training of policy and reward models. It addresses unreliable intermediate reasoning and reports stronger search and self-training performance than prior approaches.

  • ReST-MCTS* self-trains policy and process reward models using high-quality samples generated by reward-guided tree search.
  • ReST-MCTS* outperforms other self-training paradigms and achieves higher accuracy than previous reasoning baselines under the same search budget.
  • Process reward models provide finer-grained supervision than outcome reward models for identifying false-positive solutions.
  • ReST-MCTS* integrates process reward guidance with tree search to explore solution spaces and synthesize high-quality trajectories.

A.3 LLM Self-Training

This section describes iterative LLM self-training, value-model supervision, and the MCTS* search procedure. The approach generates and verifies traces, infers bounded quality values, and uses them to guide expansion and rollout.

  • Generation: Self-training methods generate reasoning traces and final answers, then update generators using correct samples or verifier-derived preference data.
  • Weighted Value: The weighted reward is bounded by |1 − v_k−1|, while the quality value is recursively updated as v_k = max(v_k−1 + w_sk, 0).
  • Quality Value: Quality values remain within [0, 1] under the stated initialization and reward bound.
  • Value Model Training Data: Verified search trees provide target quality values for partial solutions, forming the initial value-model dataset.
  • Value Model Training Data: The initial value-model training set can be constructed more cheaply than PRM datasets requiring human annotation.
  • Node Selection: MCTS* selects children with UCB, balancing quality value and visit count to explore promising and underexplored nodes.
  • Thought Expansion: A node with quality value at least 0.9 is returned directly as a final solution, while greedy Monte Carlo rollouts provide additional value-estimation information.

C.2 Data Generation Process and Specific Example for Reward Inference

ReST-MCTS* generates training data through search, pruning, verification, and reward inference. Its value design incorporates process reward and reasoning distance, and its iterative procedure refines both policy and critic models.

  • Data Generation: The self-training data-generation pipeline consists of search, prune, verify, and reward-inference stages.
  • MCTS* Design: Unlike AlphaLLM, ReST-MCTS* treats reasoning steps as actions and uses one value model to evaluate intermediate nodes.
  • Reward and Value: Its quality value combines process reward with reasoning distance, shaping critic training, data synthesis, and filtering.
  • Self-Training: ReST-MCTS* updates the value model across self-training iterations, whereas AlphaLLM keeps its critic static and focuses mainly on policy improvement.

E.1 Training and Evaluation of Initial Value Model

The initial value model is fine-tuned to predict correctness and completeness for partial solutions. Evaluation on a 14k-sample test set reports 69.3% accuracy within the stated tolerance.

  • Training: ChatGLM3-6B and Mistral-7B are fine-tuned on split value data to predict partial-solution values.
  • Training: A linear layer transforms model probabilities into scalar values, optimized with AdamW and MSE loss at learning rate 1e-6.
  • Evaluation: The value model is evaluated on a 14k-sample test set using an absolute tolerance of 0.1.
  • Evaluation: 69.3% accuracy is achieved by the initial value model, indicating reliability in most evaluated situations.

E.2 Benchmark Setup

The benchmark compares search and verification methods under controlled token budgets, using standardized labeled science and math tasks. ReST-MCTS* uses a PRM-guided tree search whose cost and accuracy are evaluated against several baselines.

  • The benchmark evaluates CoT, ToT, and ReST-MCTS* on labeled datasets including MATH, SciBench, and SciEval.
  • Verification baselines: ORM + BoN samples N solutions and selects the one with the highest outcome-reward score.
  • ReST-MCTS*: ReST-MCTS* uses Vθ as a PRM to guide MCTS*, with token usage controlled by iteration number T and branch parameter b.
  • Verification baselines: Self-Consistency selects the most frequent final answer among N chain-of-thought solutions.
  • Verification baselines: PRM + BoN uses DFS-based search and selects the solution with the highest product of step-wise PRM scores.
  • Efficiency results: MCTS* generally requires fewer tokens than Self-Consistency and ORM + BoN to reach a given expected accuracy on MATH and SciBench.
  • Efficiency results: MCTS* spends more time exploring and simulating, but its value design avoids massive Monte Carlo estimation and keeps time consumption reasonable.

E.5 ReST-MCTS∗on SciEval

On SciEval, ReST-MCTS* achieves the highest overall accuracy for both evaluated model backbones, although tree search is less advantageous on easier, short single-choice questions. The section also situates these comparisons within broader reasoning and MCTS procedures.

  • SciEval results: 79.87% and 62.31% overall accuracy are reported for ReST-MCTS* with GLM4 and GPT-3.5-turbo, respectively, exceeding other baselines.
  • SciEval results: Tree-search methods do not improve CoT on some SciEval subsets that mainly contain easier single-choice questions requiring few reasoning steps.
  • Search procedure: Classical MCTS balances exploration and exploitation through UCT, combining average reward with an exploration term based on visit counts.

H Limitations

The paper’s limitations concern generalization, value-model scale and diversity, and training-data filtering. Current evidence is concentrated on mathematical and scientific reasoning with oracle labels.

  • Generalization: ReST-MCTS* still relies on ground-truth oracle labels and has not yet demonstrated generalization to coding, agents, conversation, or other unlabeled tasks.
  • Generalization: Tasks requiring complex multistep planning without ground-truth answers need better reward feedback and more generalizable reward models.
  • Value-model scale and diversity: The proposed value model would benefit from a larger backbone and more diverse training data beyond mathematical and scientific reasoning.
  • Data filtering: The authors identify stronger and larger value-model backbones as a possible way to improve filtering of high-process-value trajectories.
  • Scope: The method aims to automate process-reward labeling and generate higher-quality datasets, but the paper notes remaining limitations in the approach.

J Reproducibility

The paper reports reproducibility support through documented training, search, evaluation, and experimental details, together with released implementation code. Reviewer responses point to dedicated sections and repository access for verification.

  • Implementation: The authors state that training code, the tree-search algorithm, and evaluation details are public in their repository.
  • Documentation: Training details for the value model, self-training backbones, and experimental settings are provided in Section 4.1.
  • Implementation: The enhanced MCTS* algorithm is documented through Algorithm 1 and the public code.
  • Documentation: Evaluation details covering iterative self-training, value models, search-budget comparisons, and reasoning policies are provided in Sections 4.2 and 4.3.
  • Implementation: The submission states that all research code was uploaded to the ReST-MCTS repository.

7. Experiment Statistical Significance

The paper’s review materials address statistical reporting, reproducibility resources, ethics, societal impacts, asset documentation, and human-subjects requirements.

  • Statistical significance: The authors state that experimental statistical significance is provided in Figure 2, Section 4.3.
  • Reproducibility: The authors state that experiment compute resources are provided in Table 2, Section 4.
  • Ethics and impacts: The authors state that the work conforms to the NeurIPS Code of Ethics and uses open-source models and datasets with citations.
  • Ethics and impacts: The authors state that broader impacts are provided in Section I.
  • Assets and documentation: The authors state that original assets are credited, proper licensing is selected, and code is submitted in an anonymized zip file.
  • Human subjects: The authors state that the research involves neither crowdsourcing nor human subjects, so participant-instruction and IRB requirements are marked not applicable.
Loading 2406.03816v3…