PaperlayerView exact revision on arXiv

Source-linked AI summary

Enhancing LLM Reasoning with Reward-guided Tree Search

Jinhao Jiang, Zhipeng Chen, Yingqian Min, Jie Chen, Xiaoxue Cheng, Jiapeng Wang, Yiru Tang, Haoxiang Sun, Jia Deng, Wayne Xin Zhao, Zheng Liu, Dong Yan, Jian Xie, Zhongyuan Wang, Ji-Rong Wen

arXiv:2411.11694v4cs.CLcs.AI

TL;DR

Complex reasoning remains a challenge for LLMs, while reproducing o1-like reasoning systems entails substantial implementation choices. The paper presents STILL-1, a reward-guided tree-search framework, and reports improved performance on four mathematical benchmarks. Its results are preliminary, and the authors say the current implementation remains far from their ideal reasoning system.

  • Problem

    LLMs remain limited on complex reasoning tasks, and reproducing o1-like systems is not straightforward.

    Source: §1 · pp. 1–2

  • Method

    STILL-1 integrates policy and reward models with tree search, with the policy navigating a dynamically expanding tree under reward-model guidance.

    Source: §1 · p. 2

  • Results

    STILL-1 significantly enhances policy-model performance on MATH-OAI, GSM-Hard, Olympiad Bench, and College Math.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The report documents design choices for implementing reward-guided tree search, including policy training, reward modeling, and search.

    Source: §1 · p. 2

    The authors describe the work as preliminary and say the current implementation is far from an ideal system that is scalable, simple, capable, and general.

    Source: §4 · p. 17

Abstract

from arXiv · show

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. The implemented framework is denoted as \textbf{STILL-1}. 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.

1 Introduction

The paper presents STILL-1 as a reward-guided tree-search framework and reports improved policy-model performance on four challenging mathematical benchmarks. It also documents design choices across policy training, reward modeling, and search to inform implementation of search-based reasoning systems.

1 Introduction

  • Search-based reasoning frameworks combine a policy model that generates reasoning steps, a reward model that guides them, and a search algorithm that directs exploration.The paper notes that these approaches show promise but have not reached o1's performance.Source: §1 · p. 2
  • STILL-1 uses a policy model to navigate a dynamically expanding tree guided by a specially trained reward model.The report describes implementation choices for this reward-guided tree-search framework.Source: §1 · p. 2
  • The policy-model investigation covers adapting reasoning to the search-tree format and using reward-guided preference optimization.Source: §1 · p. 2
  • The reward-model investigation compares discriminative and generative forms, supervision types, and ranking- versus score-based optimization, including iterative refinement with the policy model.Source: §1 · p. 2
  • The search component implements an MCTS-like algorithm and explores changes intended to improve its effectiveness and efficiency for mathematical reasoning.Source: §1 · p. 2
  • STILL-1 significantly enhances policy-model performance on MATH-OAI, GSM-Hard, Olympiad Bench, and College Math.The authors also report an empirical analysis of policy, reward-model, and tree-search design.Source: §1 · p. 3

2 Method

For mathematical problems described in text, STILL-1 combines a policy model, reward model, and tree-search algorithm. The policy generates reasoning steps along a search-tree path, while search structures the process and reward feedback guides actions and exploration.

2 Method

  • STILL-1 combines a policy model, reward model, and search algorithm to address mathematical reasoning tasks described in text.Source: §2 · p. 3
  • The policy model generates reasoning steps from a partial solution prefix along a search-tree path, while the search algorithm constructs the tree and the reward model supplies guidance.Source: §2 · p. 3
  • The framework trades test time for improved accuracy by enabling broader, more informed exploration of the reasoning space.Source: §2 · p. 3
  • Figure 1 presents an overview of the STILL-1 reasoning framework.Source: Figure 1 · p. 3

2.1 Policy Model

The policy-model process adapts the model to a structured reasoning format before preference optimization. The format treats each complete logical problem-solving step as a tree node, and reward-model feedback is used to construct preference data and improve the policy.

2.1 Policy Model

  • Policy-model training consists of instruction tuning for reasoning-format adaptation followed by preference optimization for policy improvement.Source: §2.1 · p. 3

2.1.1 Instruction Tuning for Reasoning Format Adaptation

  • A reasoning-tree node represents one complete logical step, which may span multiple sentences in a mathematical solution.The format organizes a response as a problem rephrasing, labeled steps with detailed calculations or methods, and a final answer.Source: §2.1.1 · p. 3
  • Formatted solutions are generated from NuminaMath with Qwen2.5-Math-72B-Instruct, filtered for correctness, and used to instruction-tune the policy model.Source: §2.1.1 · p. 4

2.1.2 Preference Optimization for Policy Improvement

  • Preference data are built by sampling diverse solutions, removing low-quality outputs, labeling correctness, and pairing high-ranked correct and incorrect samples scored by the reward model.The method keeps one positive-negative pair per problem to increase training-data diversity and prevent overfitting.Source: §2.1.2 · pp. 4–5
  • The policy is optimized on the constructed preference pairs using direct preference optimization.The paper defines π as the generative probability of a sample text given its question and model, and Θref as the reference-model parameters.Source: §2.1.2 · p. 5

2.2 Reward Model

The reward model is designed to guide policy reasoning in the search framework. The paper describes its design choices, training data and procedures, and iterative refinement with the policy model.

2.2.1 Key Design Considerations in Reward Modeling

  • The reward model is a language model that supplies feedback signals to guide policy reasoning, and the design considers discriminative versus generative, outcome versus process supervision, and ranking versus scoring objectives.The authors choose scoring-based optimization because the search framework uses concrete scores for solution evaluation.Source: §2.2.1 · pp. 5–6
  • The paper adopts a generative, outcome-supervised, scoring-based reward model, with this configuration informed by experiments and available computational resources rather than claimed as optimal.The model assesses whole-solution correctness; the authors also report that outcome-supervised training can provide some step-level assessment.Source: §2.2.1 · pp. 5–6, §2.2.3 · p. 6

2.2.2 Training Data Construction

  • Reward-model training data pairs policy-generated candidate solutions with ground-truth correctness labels, then removes redundant responses and balances correct and incorrect examples per problem.The model is trained to provide outcome-level absolute scores; active learning selects high-score correct and incorrect responses for further training.Source: §2.2.2 · p. 6, §2.2.3 · pp. 6–7

2.2.3 Reward Model Training

  • The reward-model backbone is adapted to mathematical tasks by fine-tuning on filtered solutions generated for sampled NuminaMath problems.The selected backbone is LLaMA-3.1-8B-Instruct; the paper describes sampling 70K problems and removing incorrect generated solutions.Source: §2.2.3 · p. 6

2.2.4 Iterative Training for Mutual Evolution

  • The framework iteratively refines policy and reward models, using policy-generated solutions to train the reward model and its scores to preference-optimize the policy.The resulting models serve as the backbone for the next iteration.Source: §2.2.4 · pp. 7–8

2.3 Search Algorithm

The search procedure uses policy-generated continuations and reward-model guidance to explore a question-dependent tree. The paper describes selection, expansion, simulation, and backpropagation, then presents pre-expansion and reward-evaluation enhancements.

2.3.1 Overall Process

  • MCTS selection traverses from the root by choosing the child with the highest UCB, balancing reward value against visit frequency; MCTSG instead selects all leaves whose reward exceeds a threshold based on leaf-value mean and standard deviation.The threshold is dynamically calculated, with λ controlling how far it exceeds the mean.Source: §2.3.1 · p. 8
  • When a selected leaf is nonterminal, the policy expands it into k children, and rollouts from each child are scored by the reward model and averaged to initialize its value.The initial reward values are propagated back through the tree, updating node statistics such as visit counts and state values.Source: §2.3.1 · pp. 8–9
  • Pre-expansion opens many nodes in the initial tree layers before ordinary search, increasing early coverage of possible actions when reward-model discrimination is limited.The paper describes this as expanding all nodes in the initial layers rather than only the selected promising leaf.Source: §2.3.1 · p. 9
  • MCTS, MCTSG, and beam search offer different trade-offs: local versus all-leaf selection, potential exploration of incorrect paths, and parallelizable but limited exploration, respectively.The authors note that MCTS can make suboptimal local selections, while MCTSG may spend computation on incorrect paths.Source: §2.3.1 · p. 9

2.3.2 Performance Optimization

  • Rollout samples also provide historical answer frequencies for a self-consistency score, which the framework combines with reward-model scores; the authors report that self-consistency often predicts ground-truth answers accurately, particularly on relatively easy questions.The combined reward weights the self-consistency score by α.Source: §2.3.2 · p. 9
  • A calculator verifies equations in generated steps with SymPy and returns feedback on whether computed results match, addressing calculation errors that reward models and self-consistency checks may miss.The paper motivates this tool by noting that small miscalculations can accumulate and disrupt a solution.Source: §2.3.2 · p. 10

3 Experiments

This section reports experiments examining the effectiveness of the implemented framework.

3 Experiments

  • The experiments examine the effectiveness of the implemented framework.The supplied passage states the purpose of the section but provides no experimental outcomes.Source: §3 · p. 10
  • The section introduces an experimental evaluation of the framework.No benchmark or result details are supplied in this passage.Source: §3 · p. 10
  • The reported experimental work is intended to assess the implemented framework.The passage does not specify the evaluation setup or findings.Source: §3 · p. 10

3.1 Evaluation Setup

The evaluation uses LLaMA-3.1-8B-Instruct as the policy and reward backbone and reports average performance across four mathematical benchmarks. Three benchmarks are evaluated on random 500-sample subsets to save testing time.

3.1 Evaluation Setup

  • The evaluation covers MATH-OAI, GSM-Hard, OlympiadBench, and College Math, with full test sets of 500, 1319, 675, and 2818 examples, respectively.For the latter three benchmarks, the study randomly samples 500 examples each to save testing time.Source: §3.1 · p. 10
  • LLaMA-3.1-8B-Instruct serves as the backbone for both policy and reward models because the authors report it has strong overall capabilities without saturating these benchmarks.Source: §3.1 · p. 10
  • The study uses the evaluation tools employed in prior research and reports each method's average performance over all test problems.Except for the main experiments, reported results primarily use MATH-OAI unless otherwise specified.Source: §3.1 · p. 10

3.2 Main Results

STILL-1 achieves the strongest overall performance among the compared methods across four benchmarks. In iterative training, reward-based selection outperforms random selection, and policy and reward models improve across two iterations on almost all datasets.

3.2.1 Overall Performance Comparison

  • STILL-1 improves the reference baseline by 46.9%, 7.3%, 91.6%, and 31.4% across the four benchmarks, achieving the best performance among compared methods.Table 1 reports the performance comparison and gains over the baseline.Source: §3.2.1 · p. 10, Table 1 · p. 11

3.2.2 Results of Iterative Training

  • Reward-based selection outperforms random selection for both models, and policy and reward models improve across two iterations on almost all datasets.The authors attribute benefits to reward feedback for both reward-model active learning and policy-model preference optimization.Source: §3.2.2 · pp. 10–11, Table 2 · p. 11

3.3 Further Analysis of Policy Model Training

Policy-model analyses find that stronger synthesis models improve performance, while increasing synthetic-data volume yields diminishing or declining gains. Policy optimization also improves self-consistency, though it remains below the full STILL-1 framework.

3.3 Further Analysis of Policy Model Training

  • The policy-model evaluation uses accuracy for direct generation and maj@10 and pass@10 as coverage metrics over ten generated solutions.maj@10 uses majority vote; pass@10 measures recall among generated solutions.Source: §3.3 · p. 11

3.3.1 Effect of Data Synthesis Model

  • Stronger synthesis models improve policy performance, while self-generated data offers little gain and may degrade it.The authors say higher-quality solution-process data is necessary to further improve a strong foundational model; larger general-model data helps comparably to domain-specific-model data.Source: §3.3.1 · pp. 11–12, Table 3 · p. 12

3.3.2 Effect of Training Data Scaling

  • More synthetic data generally improves reasoning-format adaptation, but gains may taper or decline, while pass@k shows no significant improvement.The experiments compare 5K, 10K, 20K, 30K, and 40K synthetic examples; the framework uses 10K to balance synthesis cost and downstream performance.Source: Figure 2 · p. 12, §3.3.2 · p. 12

3.3.3 Performance Improvement on Self-consistency

  • Policy optimization significantly improves self-consistency across sampling budgets, but enhanced-policy self-consistency remains below the full STILL-1 framework.The sampling budget is the number of rollouts; the original policy serves as the unoptimized comparison.Source: §3.3.3 · pp. 12–13, Figure 3 · p. 13

3.4 Further Analysis of Reward Model Training

The reward-model analysis finds that domain adaptation and careful data selection improve reward quality, while model design and training choices affect assessment performance. An outcome-supervised reward model also provides some process-level supervision in search tasks.

3.4.1 Effect of Model Adaptation Strategies

  • Mathematical-instruction fine-tuning generally improves reward-model accuracy as BoN-N increases, while format adaptation helps at N = 5 and 10 but is not consistently beneficial as N grows.The reported adaptation comparisons appear in Table 4.Source: §3.4.1 · p. 13

3.4.2 Effect of Data Selection Strategies

  • Deduplication and balancing positive and negative responses improve reward-model performance, with removing the balancing strategy causing the larger drop.The authors suggest models may overfit to shortcuts in biased datasets.Source: §3.4.2 · p. 14, Table 5 · p. 14
  • Ranking positive and negative samples by descending score performs best among the tested selection methods and significantly improves on random sampling.Selecting positive samples in ascending score order significantly reduces reward-model effectiveness; high-scoring negative samples are more beneficial.Source: §3.4.2 · p. 14, Table 6 · p. 14

3.4.3 Effect of Model Design and Training

  • Binary cross-entropy L1 outperforms the other tested objectives, while ranking-based L4 cannot provide absolute scores and may not suit MCTS.L3 and L4 perform comparably with fewer candidates, while L4 performs well at BoN-20.Source: §3.4.3 · p. 14, Table 7 · p. 15
  • Generative reward models outperform discriminative models, larger backbones improve assessment, and the math-specific model outperforms the large-scale general model.The comparison covers discriminative and generative reward models.Source: §3.4.3 · p. 15, Table 8 · p. 15

3.4.4 Performance of Reward Model on Process-supervised Tasks

  • Outcome-supervised reward models outperform the CoT baseline on depth-first and beam-search tasks, with larger gains for beam search.The results suggest outcome-trained reward models can provide some process-level supervision; beam search retains more candidates at each step.Source: §3.4.4 · p. 15

3.5 Further Analysis of Search Algorithm

The search analysis compares three algorithms and tests pre-expansion. MCTSG has the highest accuracy among those tested, and two-layer pre-expansion balances search efficiency and computational cost.

3.5.1 Comparison of Different Search Algorithms

  • MCTSG achieves the highest accuracy among the tested methods, while beam search has the lowest; the comparison is on MATH-OAI.Figure 4 presents search-algorithm performance in panel (a).Source: §3.5.1 · p. 16, Figure 4 · p. 16
  • MCTS uses UCB-based expansion to balance exploration and exploitation, while beam search's forward-only expansion can accumulate uncorrected errors.Extensive MCTS exploration can sometimes slightly decrease accuracy.Source: §3.5.1 · p. 16

3.5.2 Effect of Pre-expansion

  • Pre-expansion improves accuracy over no pre-expansion, and two layers offer an overall balance between computational cost and search efficiency.The tested range is zero to two pre-expansion layers; Figure 4 panel (b) presents the comparison.Source: §3.5.2 · p. 16, Figure 4 · p. 16

4 Conclusion

The conclusion presents STILL-1 as a preliminary search-enhanced reasoning framework and reports effectiveness on mathematical reasoning tasks. The authors also describe their implementation as far from the desired scalable, simple, capable, and general system.

4 Conclusion

  • STILL-1 integrates policy, reward, and search components, and preliminary results demonstrate effectiveness on mathematical reasoning tasks.The paper characterizes this work as a technical report on preliminary exploration.Source: §4 · p. 17
  • The team reports difficulty configuring design factors, integrating components, and jointly optimizing the framework.The authors say they are still searching for the right research pathways to reproduce o1-like systems.Source: §4 · p. 17
  • The current implementation remains far from the authors' desired system, which they describe as scalable, simple, capable, and general.The authors plan to continue the project in future work.Source: §4 · p. 17

Related papers

Loading 2411.11694v4…