Boosting Policy and Process Reward Models with M onte C arlo Tree Search in Open-Domain QA (2025)

Abstract: The recent introduction of OpenAI’s O1/O3 model represents a significant milestone in developing strong reasoning capabilities in Large Language Models (LLMs). By introducing more computational budget during test-time, LLMs have the potential to explore more accurate and higher-quality solutions. However, such paradigms are primarily verified in domains that have well-defined criteria for responses, such as coding and mathematics. Inspired by the success of this paradigm, we aim to bridge it to more subtle open-domain question answering. Specifically, we utilize search mechanisms such as Monte Carlo Tree Search (MCTS) for both policy model improvement and reward model improvement that achieve better performance in test-time scaling strategies. Our contributions are summarized in two folds: For the training phase, we demonstrate that our approach surpasses previous SOTA automatic data annotation methods and various public instruction-tuning datasets, with fewer data points. This offers a more data-efficient solution for training robust models. For the inference phase, we utilize the intermediate values collected during training data construction to train a process reward model called PRM+. This model employs a novel two-stage training method to provide finer-grained guidance across the generation trajectory. This introduces no additional overhead during training data collection and further enhances performance by scaling test-time computation. Experimental results show that our method can effectively improve the performance of both the policy model and the reward model.

## introduction large language models (llms) have made remarkable progress in recent years, demonstrating strong reasoning capabilities across a wide range * co-corresponding author † corresponding author of tasks @xcite @xcite . a key development in this area is the ability of llms to generate long chains of thought (cot) @xcite . this enables them to iteratively refine their outputs and improve response quality @xcite @xcite . the process can be viewed as a search over the space of possible reasoning trajectories, where increased computational resources during inference allow llms to explore more accurate and higher-quality solutions @xcite . in addition to sequential self-refinement of their own outputs, parallel decoding is another effective strategy for scaling test-time computation. previous work suggests that by simply applying repeated sampling to the same question @xcite , models can achieve relatively high performance. moreover, when combined with a welllearned value function @xcite @xcite ) that provides finegrained step-level supervision, models can navigate higher-quality reasoning paths to solve complex problems. although the above-mentioned methods have already proven highly effective in structured domains such as mathematical reasoning @xcite @xcite and code generation @xcite @xcite where there are well-defined ground truths, their applicability in open-domain tasks remains underexplored. unlike structured problems, open-domain question answering (qa) lacks explicit correctness criteria, making it challenging to systematically guide the search process toward optimal responses. to address this gap, our work explores whether search-based methods can enhance open-domain qa in both training and inference phases. we hypothesize that guiding llms through a structured init_max selection 0. initial rollout b backpropagation s simulation c expansion intermediate node expanded node terminal node select backup rollout repeat until better solution found (@xmath0.3 0.8 0. 9 0.9 cur_max 1.2 0.8 search process can help identify high-quality responses, leading to more robust model training and better inference-time generation. specifically, we propose leveraging monte carlo tree search (mcts) @xcite to explore multiple generation paths and select those that maximize a reward model's evaluation score. by treating generation as a search problem, we systematically improve the model's ability to generate informative and well-grounded answers. our approach comprises two key phases: 1) the improvement of the policy model and 2) the improvement of the reward model. in the training phase, we employ mcts to collect high-quality trajectories by selecting the bestscoring generation paths during searching. these paths are then used to construct a supervised fine-tuning (sft) dataset. additionally, we create a direct preference optimization (dpo) dataset by pairing the highest and lowest-scoring responses. this allows the model to learn preference-based refinements directly from the data. compared to previous automatic data annotation methods and publicly available instruction-tuning datasets @xcite @xcite , our approach identifies superior response trajectories and achieves better downstream task performance using fewer data points. this leads to a more data-efficient training process overall. during the inference phase, we further enhance the policy model's performance by introducing a process reward model (prm) to facilitate scalable test-time computation. unlike conventional outcome reward models (orm), which evaluate only the final output, the prm captures incremental improvements in generation quality and provides a more structured mechanism for refining responses at each step. notably, collecting data to train the prm does not introduce additional computational overhead, as the intermediate values used as supervision signals are by-products of the training data collection process. furthermore, we propose a two-stage training method to develop an enhanced version of the prm, referred to as prm+ in our paper. experimental results demonstrate that our prm+ outperforms various orms in scaled settings, validating the effectiveness of process-based supervision. our key contributions are as follows: process reward model one of the key components for effectively scaling inference-time computation is the development of a reward model capable of evaluating the quality of generated candidates @xcite . orms are widely adopted for this purpose. however, a primary challenge with orms is their coarse-grained feedback, which can lead to situ-ations where, despite the final answer being correct, the solution paths may contain flaws. to address these limitations, process reward models have recently been developed. currently, there are two primary annotation methods: 1) human annotation, which requires manually rating intermediate steps of the solution paths, making it costintensive and time-consuming @xcite ; and 2) automatic annotation, which includes methods such as monte carlo estimation to construct stepwise supervision signals @xcite , leveraging llm-as-a-judge to generate signals @xcite , or deriving signals from outcomebased feedback @xcite . despite the significant progress and potential of these methods, they have primarily been validated in mathematical domains. in this work, we aim to explore the effectiveness of this paradigm in open-domain settings, which present more subtle optimization challenges. another concurrent work with a similar goal is @xcite , which utilizes segmentation and aggregation methods to curate intermediate signals and focuses solely on training a prm. ## methodology in this section, we first present the problem formulation, notation, and primary objective of our paper in section 3.1. following this, we detail our methodology for curating sft datasets and dpo datasets to optimize the policy model in section 3.2. afterwards, we outline the data curation and training process for our prm+, which is designed to explore inference-time scaling, in section 3.3. we defer the details of the search algorithm employed in our work to appendix a. ## conclusion in this paper, we introduce a search-based method to enhance both policy and reward models for opendomain qa. our approach achieves data-efficient training, outperforming previous sota data synthetic methods while requiring fewer data points. additionally, we propose a two-stage training strategy to build prm+ which demonstrates a superior scaling trend as computation increases during inference.