https://arxiv.org/abs/2411.11694
Technical Report: Enhancing LLM Reasoning with Reward-guided Tree Search
Recently, test-time scaling has garnered significant attention from the research community, largely due to the substantial advancements of the o1 model released by OpenAI. By allocating more computational resources during the inference phase, large languag
arxiv.org
Recently, test-time scaling has garnered significant attention from the research community, largely due to the substantial advancements of the o1 model released by OpenAI.
By allocating more computational resources during the inference phase, large language models (LLMs) can extensively explore the solution space by generating more thought tokens or diverse solutions, thereby producing more accurate responses.
However, developing an o1-like reasoning approach is challenging, and researchers have been making various attempts to advance this open area of research.
In this paper, we present a preliminary exploration into enhancing the reasoning abilities of LLMs through reward-guided tree search algorithms.
This framework is implemented by integrating the policy model, reward model, and search algorithm.
It is primarily constructed around a tree search algorithm, where the policy model navigates a dynamically expanding tree guided by a specially trained reward model.
We thoroughly explore various design considerations necessary for implementing this framework and provide a detailed report of the technical aspects.
To assess the effectiveness of our approach, we focus on mathematical reasoning tasks and conduct extensive evaluations on four challenging datasets, significantly enhancing the reasoning abilities of LLMs.
2 Method
In this work, we focus on the mathematical domain, specifically addressing mathematical problems described in the text.
We implement a reward-guided tree search framework designed to enhance the reasoning capabilities of LLMs.
This framework consists of three main components: the policy model, the reward model, and the search algorithm.
Within our framework, the policy model generates new reasoning steps based on a partial solution prefix along a given path in the search tree.
The search algorithm constructs the search tree to structure the reasoning process, while the reward model provides feedback signals to guide the policy model's actions and the search process.
This approach allows the policy model to explore a broader reasoning space in a more informed manner, increasing the likelihood of finding the correct answer.
Overall, it trades test time for improved accuracy.
In the following sections, we will provide detailed descriptions of the implementation of these three key components.
2.1 Policy Model
In this section, we provide a detailed description of the training process for the policy model, which primarily consists of two steps: instruction tuning for reasoning format adaptation (Section 2.1.1), and preference optimization for policy improvement (Section 2.1.2).
2.1.1 Instruction Tuning for Reasoning Format Adaptation
As we perform the reasoning using a tree-structured search algorithm, this necessitates adapting the reasoning format of the policy model.
Defining the Reasoning Format.
To define the reasoning format, we first set the granularity of each reasoning step, i.e., how each node is represented, framing the multi-step reasoning process as a node-based search within a tree structure.
Previous studies have extensively explored reasoning steps at both the token and sentence levels [28, 29].
However, in mathematical contexts, a logically complete step in problem-solving may involve multiple sentences.
Therefore, we model a complete logical step within the mathematical problem-solving process as a node in the tree.
Below is the output template for our reasoning process.
Step-by-Step Output Format in Our Reasoning Process
-----------------------------------------------------------------
**Problem Formulation**
Rephrasing the problem in your own words to capture its essential meaning
**Step 1: Brief title that summarizes the overall goal of that step**
Detailed calculation or problem-solving process of this step.
...
**Step i: Brief title that summarizes the overall goal of that step**
Detailed calculation or problem-solving process of this step.
...
**Final Answer**
Output the final answer in boxed{}.
------------------------------------------------
Supervised Format Following.
As shown in the prompt template, our reasoning format begins by suggesting that the LLM thoroughly understands the problem, followed by a sequence of individual steps, each with a summarization label and a detailed description or calculation.
To align with this desired reasoning format, we synthesize formatted data through in-context learning with a more capable LLM, followed by instruction tuning on the policy model.
Specifically, using the NuminaMath dataset [30], we employ the Qwen2.5-Math-72B-Instruct model [31] to generate formatted solutions for each problem via one-shot in-context learning.
After filtering out incorrect solutions, we conduct supervised fine-tuning on the policy model with these instruction data to effectively improve the reasoning abilities while ensuring it adapts well to the reasoning format.
We utilize the following prompt template for our policy model during the supervised fine-tuning and later inference.
Prompt Template for Policy Model
Analyze and respond to the following question step by step.
Begin by rephrasing the problem in your own words to capture its essential meaning accurately.
Then, proceed to solve the problem systematically, ensuring that each step is introduced with a concise heading that summarizes its objective.
Follow this with detailed explanations of the calculations or methodologies employed in the problem-solving process.
Finally, present the final answer within boxed{}.
{Problem}
2.1.2 Preference Optimization for Policy Improvement
By aligning the policy model with the desired reasoning format, we can better control the step-by-step generation of the reasoning process.
This alignment facilitates both the execution of tree search algorithms and the assessment by the reward model.
Next, we conduct preference optimization on the policy model using feedback from the reward model (Section 2.2), thereby achieving policy improvement.
This process involves two key steps: constructing the training data and performing preference optimization.
Next, we introduce the implementation of these two steps.
Training Data Construction.
To improve the policy model, we construct the corresponding training data through the following three steps: first, sample multiple solutions for each question from the policy model; second, filter out low-quality solutions; and third, create positive and negative sample pairs based on scores from the reward model and the annotated labels.
Specifically, given a set of mathematical problems Q = {}
,we begin by solving each problem qi using the aforementioned policy model after adapting its reasoning format.
During the generation process, step-level diversity is crucial, as sufficiently varied paths are essential for exploring different actions within the tree.
To address this, we increase the temperature parameter to enhance sampling diversity, thereby obtaining more diverse solution paths {s˜i,j} k j=0.
Next, we apply a rule-based approach to remove samples that contain garbled content or deviate from the specified reasoning format.
We then annotate the remaining samples with correctness labels by comparing the predicted answers to the ground-truth si , and only retain those problems with both correct and incorrect samples.
Subsequently, we score each retained sample using the reward model and select the top-ranked correct and incorrect samples to form positive-negative sample pairs.
This approach aims to select highly confident positive samples and incorrect samples that are nearly correct, providing more informative signals for preference optimization.
Similar strategies have been also used in previous work [13, 14, 12, 18].
Additionally, we sample only one pair of positive and negative examples per problem to increase the diversity of the training data and prevent overfitting.
Ultimately, we obtain the training data set {⟨qi , s+ i , s− i ⟩}N i=0.
Preference Optimization.
After obtaining the training data, we apply the chat template and prompt to the training data and then conduct preference optimization on the policy model P using the direct preference optimization algorithm [11] as follows:
where σ(·) denotes the sigmoid function, Θ_ref denotes the parameters of the reference model (i.e., the policy model itself), and π denotes the generative probability of a sample text given the question and the corresponding model
Note that although we describe a single-turn policy improvement process above, it can be easily extended to multi-turn optimization using mutually optimized reward models, which will be detailed in Section 2.2.4.
2.2 Reward Model
In this section, we first identify the key design considerations in implementing the reward model, and then introduce its detailed training process.
2.2.1 Key Design Considerations in Reward Modeling
Since the reward model (RM) provides feedback signals to guide the policy model's reasoning within the search framework, it plays a crucial role in directing the reasoning process of LLMs.
Generally, a reward model can be implemented using various types of machine learning models [32, 33, 34]; however, in our work, we focus on using LLMs as the backbone model.
To design an effective RM, we concentrate on the following three key considerations.
Discriminative RM v.s. Generative RM.
When implementing reward models for LLM reasoning, we can consider using either discriminative RM or generative RM.
In existing work [35, 34, 31], discriminative RM has been widely utilized to provide supervision signals, which projects the model's hidden state into a score that serves as the supervision signal.
In contrast, generative RM takes a specific prompt as input (i.e., guidance for evaluation and the text to be evaluated) and generates a descriptive evaluation of the quality for the provided solution, potentially accompanied by associated reward scores [36, 37, 38].
A direct implementation of generative RM involves constructing assessment prompts, such as "Is the answer correct (Yes/No)?", and then using the prediction probabilities of the assessment tokens (e.g., the probability of "Yes") as the reward score.
Compared to discriminative RM, a potential advantage of generative RM is that it can leverage the learned knowledge of the base model more effectively, as its training process closely aligns with both the pre-training and post-training phases of the model.
Outcome-supervised RM v.s. Process-supervised RM.
The second issue we consider is the granularity of the supervision signals provided by the reward models.
Typically, we can use either outcome-supervised reward models (ORM), which assess the correctness of the entire solution, or process-supervised reward models (PRM), which evaluate the correctness of intermediate steps in the reasoning process [24, 35, 39].
To develop these two types of reward models, it requires labeled data at different granularities: solution-level labels and step-level labels.
In practice, solution-level labels are relatively easy to obtain when ground-truth answers are available for the tasks, whereas step-level labels are very difficult to annotate.
To automatically obtain process-supervised signals, we can employ rollout-based methods (e.g., Math-Shepherd [35]) to derive step-level labels from ground-truth answers.
The core idea is to annotate each reasoning step based on whether correct answers can be derived from it through a rollout process.
However, rollout-based methods typically require more computation time, particularly when the process is repeated multiple times.
As a result, this work focuses on training outcome-supervised reward models, primarily relying on outcome-level supervision signals to guide the reasoning process.
Interestingly, we find that the reward model trained with solution-level labels also has potential in assessing step-level correctness.
Ranking-based RM v.s. Scoring-based RM.
To train reward models, we can explore different optimization objectives, which are generally either ranking-based or scoring-based.
Specifically, a ranking-based RM is trained to identify which candidate response is the best, focusing primarily on the preference order among candidate solutions.
In contrast, a scoring-based RM assigns an absolute reward score for a single response based on its assessed quality level.
The following equations illustrate the comparison between ranking-based and scoring-based approaches:
Since our reasoning framework relies on the guidance of concrete scores for solution evaluation, we adopt the scoring-based optmization objective for training our RM.