ReST-MCTS*: LLM Self-Training via Process Reward Guided Tree Search
This work addresses the challenge of generating effective training data for LLM self-training, which is crucial for advancing reasoning capabilities in AI systems, though it is incremental as it builds on existing self-training and tree-search methods.
The paper tackles the problem of low-quality training data in LLM self-training by developing ReST-MCTS*, a reinforced self-training approach that integrates process reward guidance with tree search to collect high-quality reasoning traces and per-step values for training policy and reward models, achieving higher accuracy than prior baselines and outperforming other self-training algorithms in iterative improvements.
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.