PaperlayerView exact revision on arXiv

Source-linked AI summary

TreeRL: LLM Reinforcement Learning with On-Policy Tree Search

Zhenyu Hou, Ziniu Hu, Yujiang Li, Rui Lu, Jie Tang, Yuxiao Dong

arXiv:2506.11902v1cs.LGcs.CL

TL;DR

On-policy LLM reinforcement learning has underused tree search, despite its potential for exploration and dense process supervision, while existing offline reward-model approaches can be ineffective. TreeRL integrates uncertainty-guided on-policy tree search with RL, and experiments report superior performance over traditional multi-chain methods across challenging math and code reasoning benchmarks.

  • Problem

    Tree search remains underdeveloped in on-policy LLM reinforcement learning, while offline process reward models can provide little performance improvement.

  • Method

    TreeRL uses EPTree to branch from high-uncertainty intermediate tokens and supplies process supervision derived from on-policy search trees to reinforcement learning.

  • Results

    TreeRL outperforms traditional independent multi-chain sampling across challenging math and code reasoning benchmarks.

  • Takeaways & Limitations

    The results highlight the potential of combining tree search and process supervision with reinforcement learning to advance complex LLM reasoning.

  • Takeaways & Limitations

    The evaluation emphasizes PassRate and requires tree search to remain efficient and highly parallelizable for modern LLM inference.

Abstract

from arXiv · show

Reinforcement learning (RL) with tree search has demonstrated superior performance in traditional reasoning tasks. Compared to conventional independent chain sampling strategies with outcome supervision, tree search enables better exploration of the reasoning space and provides dense, on-policy process rewards during RL training but remains under-explored in On-Policy LLM RL. We propose TreeRL, a reinforcement learning framework that directly incorporates on-policy tree search for RL training. Our approach includes intermediate supervision and eliminates the need for a separate reward model training. Existing approaches typically train a separate process reward model, which can suffer from distribution mismatch and reward hacking. We also introduce a cost-effective tree search approach that achieves higher search efficiency under the same generation token budget by strategically branching from high-uncertainty intermediate steps rather than using random branching. Experiments on challenging math and code reasoning benchmarks demonstrate that TreeRL achieves superior performance compared to traditional ChainRL, highlighting the potential of tree search for LLM. TreeRL is open-sourced at https://github.com/THUDM/TreeRL.

1 Introduction

TreeRL addresses limitations in applying tree search to on-policy LLM reinforcement learning by combining efficient uncertainty-guided exploration with on-policy process supervision. Its experiments report improved sampling and RL performance over traditional independent multi-chain approaches.

  • On-policy LLM reinforcement learning has largely relied on independently sampled trajectories with final-answer rewards, while tree search remains underdeveloped for this setting.
  • MCTS can be less effective and efficient than independent sampling under equal inference cost, while offline process reward models provide little RL improvement.
  • TreeRL combines tree search with reinforcement learning to generate more diverse responses and provide on-policy process supervision under the same inference cost as independent sampling.
  • EPTree branches from high-uncertainty intermediate tokens and continues generation, reducing token requirements while encouraging diverse exploration; it typically forms trees in around two iterations.
  • TreeRL derives advantage-based process supervision directly from on-policy generated trees, avoiding separate reward models and reducing susceptibility to reward hacking.
  • Experiments on college-level and competition-level math and code benchmarks show TreeRL outperforming traditional independent multi-chain sampling, with gains attributed to EPTree and process supervision.

2 Preliminary

The preliminary framework describes conventional LLM reinforcement learning as sampling multiple responses and assigning scalar rewards, then using advantage signals relative to a baseline. TreeRL instead introduces entropy-guided tree search and feeds its resulting process supervision into RL.

  • For a prompt x, a policy model generates K possible responses, and a reward function assigns each prompt-response pair a scalar reward.
  • The advantage function is commonly defined as A(x, y_i) = β(r(x, y_i) − b), where b is a baseline that varies across methods.
  • TreeRL uses EPTree to guide search by token-level uncertainty rather than traditional Monte Carlo Tree Search, then integrates the search with process-supervised reinforcement learning.

3 TreeRL: Reinforcement Learning with Tree Search

TreeRL integrates entropy-guided tree search with on-policy reinforcement learning, using branching at uncertain tokens to improve exploration and process supervision. EPTree builds diverse responses efficiently, while tree-derived advantages provide intermediate training signals without a separate reward model.

  • 3.1 Entropy-Guided Tree Search: EPTree expands branches from the top-N most uncertain tokens, using entropy to target informative regions of the reasoning tree.Unlike MCTS, it continues generation from selected intermediate tokens rather than repeatedly decomposing answers into smaller steps.
  • 3.1 Entropy-Guided Tree Search: EPTree generates multiple trees in parallel and typically requires only around 2 iterations to build a diverse and informative search tree.The tree is initialized with M parallel responses before iterative forking and expansion.
  • 3.1 Entropy-Guided Tree Search: M×(T ×N×L+1) leaves result after L expansion iterations, where M is the number of trees, N the forked points, L the iterations, and T the branching factor.Each expansion adds T candidate responses at each selected forking point.
  • 3.1 Entropy-Guided Tree Search: EPTree produces more diverse responses than independent multi-chain sampling under the same inference cost, offering more varied trajectories for RL training.The diversity comparison reports approximately 2× different responses for EPTree and MCTS relative to i.i.d. multi-chain sampling.
  • 3.2 Process Supervision from Tree Search: TreeRL assigns process supervision using node values and global and local advantages that measure prospects for correctness and improvement over a parent step.Node value is estimated as the ratio of correct descendant leaves, and the process signal combines advantage components derived from the tree.
  • 3.2 Process Supervision from Tree Search: TreeRL extracts tree sequences for RL training and downweights repeated non-leaf steps by the square root of their descendant-leaf count.This reweighting is reported to improve performance in the experiments.

4 Experiment

Experiments compare TreeRL and its EPTree sampler with multi-chain sampling and MCTS across math and code reasoning benchmarks. Results indicate improved efficiency and performance, while RL gains over SFT remain minor in the selected setting.

  • Setup: TreeRL is evaluated on six challenging reasoning benchmarks spanning MATH, Omni-MATH, AIME2024, AMC, OlympiadBench, and LiveCodeBench.Evaluation uses greedy sampling, with repeated evaluations for several datasets to reduce variance.
  • EPTree results: EPTree outperforms multi-chain sampling and MCTS on Omni-MATH-500 under the same inference cost.EPTree remains stronger across different generation costs, with its advantage over MCTS widening as inference cost increases.
  • EPTree results: EPTree exceeds multi-chain sampling by around 3% in PassRate on Omni-MATH-500.The comparison concerns different sampling methods under matched inference-cost conditions.
  • RL training: TreeRL begins outperforming ChainRL around 100 training steps and continues improving consistently across the reported training evaluations.Figure 6 reports average performance across six datasets, OlympiadBench, and Omni-MATH-500.
  • RL training: Both RL methods show only minor improvements over the SFT baseline in the selected training-data setting.The passage attributes this possibility to the RL data being based on Qwen-2.5-14B and therefore easier for the evaluated smaller model.
  • EPTree ablation: Entropy-based EPTree forking achieves better PassRate than random forking with fewer generation tokens, while both tree strategies outperform multi-chain sampling.The ablation compares branching strategies on Omni-MATH-500.
  • TreeRL ablation: Reweighted local and global advantage produces the best process-supervision performance, and removing either component reduces performance.Using all sampled responses improves more than using a subset, indicating a training-cost trade-off within the same inference budget.

5 Related Work

Prior work applies reinforcement learning and process supervision to LLMs, while tree search has mainly supported inference, data synthesis, and offline preference training. TreeRL is positioned within the less-developed use of tree search for on-policy LLM RL.

  • Reinforcement learning and process supervision: Most existing LLM RL methods use whole-response rewards, while process supervision has shown promising performance and commonly relies on a separately trained process reward model.The related work passage distinguishes outcome-level rewards from step-level process supervision.
  • Tree search for LLMs: Tree search has mainly been used for LLM alignment and inference, including preference-pair generation, SFT data generation, and process-supervision data production.These approaches generally use MCTS or reward-guided search for offline or inference-time workflows.
  • Process-signal design: The cited ablation setup distinguishes global and local advantages and records the number of responses used for training.GA and LA denote global and local advantages, respectively.

6 Conclusion

TreeRL combines on-policy tree search with process supervision for LLM reasoning. Its EPTree sampler improves response diversity and performance over MCTS and independent multi-chain sampling, and experiments report gains on math reasoning tasks.

  • Conclusion: TreeRL combines tree search with process supervision to enhance LLM reasoning.The framework applies reinforcement learning with EPTree and process supervision derived from the tree search.
  • Conclusion: EPTree improves response diversity and performance over MCTS and i.i.d. multi-chain sampling.The conclusion presents this comparison as part of TreeRL’s overall contribution.
  • Conclusion: Experiments on math reasoning tasks show that TreeRL outperforms existing techniques.The authors frame this result as evidence of the potential of RL with tree search for complex reasoning tasks.

7 Limitation

The authors identify inference-engine support and process-signal design as limitations of the current TreeRL approach. EPTree requires multiple iterations and is slower than multi-chain sampling, while reward construction remains open for improvement.

  • Inference efficiency: Current LLM inference engines lack special tree-search optimizations, leaving EPTree approximately 2× slower than multi-chain sampling.The proposed EPTree still requires 2+ iterations.
  • Process supervision: The appropriate weighting of different reasoning steps, meaningful tree-based process signals, and step-level reward normalization remain open problems.These are identified as directions requiring further exploration.

A Position Distribution Analysis

EPTree selects forking tokens by their relative position within branches, and the observed distribution is roughly uniform on Omni-MATH-500.

  • A Position Distribution Analysis: EPTree analyzes each forking point as its position divided by the branch length.The resulting relative-position distribution is plotted for Omni-MATH-500.
  • A Position Distribution Analysis: The observed forking-position distribution is roughly uniform across branches.This pattern supports the assumptions used in the theoretical analysis.

B Theoritical Analysis of EPTree

The theoretical analysis models EPTree’s branching under uniform forking positions and fixed generation lengths, showing bounded leaf-growth behavior and efficiency relative to multi-sampling.

  • B Theoritical Analysis of EPTree: Under uniformly distributed forking tokens and fixed, nonrepeating generation lengths with l ≤2, EPTree’s leaf count is theoretically bounded.The theorem assumes forking positions follow U(0, 1) within each branch.
  • B Theoritical Analysis of EPTree: The proof’s branching process uses i.i.d. uniform variables x1, x2, …, xn to represent forking positions.Figure 9 illustrates the branching configuration for n = 3 and t = 2.
  • B Theoritical Analysis of EPTree: For l = 1, the entropy-tree has Ntree = 1 + nt leaves.The proof derives this count from unit completion lengths.
  • B Theoritical Analysis of EPTree: The l = 2 case analyzes whether the next top-n entropy token lies on the main chain or a branch, then computes expected completion length.These probabilities determine the subsequent leaf-count analysis.
  • B Theoritical Analysis of EPTree: Monte Carlo simulation finds that φ decreases monotonically with n.The supplied analysis also considers the asymptotic behavior as n grows.
  • B Theoritical Analysis of EPTree: The leaf-count ratio relative to multi-sampling is shown to increase monotonically with nt.The analysis derives the ratio after comparing tree and multi-sample leaves at the same completion length.

C Detailed Evaluation of EPTree

EPTree is evaluated against multi-chain sampling and MCTS under comparable inference costs, showing improved efficiency and effectiveness in Omni-MATH-500 and stronger RL results across reasoning benchmarks.

  • C Detailed Evaluation of EPTree: 56.9 versus 52.4 is achieved by the selected EPTree RL configuration compared with chain-16.The configuration uses one interaction while maintaining a similar inference cost.

E TreeRL on General tasks

On general benchmarks, TreeRL performs comparably to ChainRL, indicating that its reported advantage in reasoning tasks does not reduce performance across diverse general-task evaluations.

  • E TreeRL on General tasks: The general-task evaluation covers MMLU-Pro, Arena-Hard, and IFEval.These benchmarks test challenging knowledge reasoning, human-like preferences, and instruction following.
  • E TreeRL on General tasks: TreeRL and ChainRL show comparable performance on MMLU-Pro, Arena-Hard, and IFEval.The evaluations use PassRate, win rate against GPT-4-0314, and the strict prompt metric, respectively.
  • E TreeRL on General tasks: The reported general-benchmark results indicate robust performance across diverse task types while TreeRL retains a particular advantage in reasoning tasks.The reasoning-task comparison is illustrated on MATH500, AMC, LiveCodeBench, and AIME2024.
Loading 2506.11902v1…