PaperlayerView exact revision on arXiv

Source-linked AI summary

TreePO: Bridging the Gap of Policy Optimization and Efficacy and Inference Efficiency with Heuristic Tree-based Modeling

Yizhi Li, Qingshui Gu, Zhoufutu Wen, Ziniu Li, Tianshun Xing, Shuyue Guo, Tianyu Zheng, Xin Zhou, Xingwei Qu, Wangchunshu Zhou, Zheng Zhang, Wei Shen, Qian Liu, Chenghua Lin, Jian Yang, Ge Zhang, Wenhao Huang

arXiv:2508.17445v1cs.LGcs.CL

TL;DR

With reasoning rewards sparse and delayed, TreePO asks how language models can explore potentially correct reasoning paths while controlling compute. It replaces independent rollouts with tree-based sampling and hierarchical advantage estimation, reporting lower reasoning costs, improved training stability, and strong performance.

  • Problem

    With rewards sparse and delayed, the paper asks how language models can explore potentially correct reasoning paths while maintaining or reducing compute.

    Source: §1 · p. 2

  • Method

    TreePO replaces independent rollouts with heuristic tree-based sampling and estimates advantages over entire subtrees to support more granular credit assignment.

    Source: §1 · p. 2

  • Results

    TreePO reduces reasoning computational costs while improving training stability and maintaining strong performance, with experiments reporting a favorable trade-off between computational cost and model performance.

    Source: §1 · p. 3, §6 · p. 14

  • Takeaways & Limitations

    The authors identify TreePO’s efficiency and structural modeling as promising for scaling reinforcement learning to more complex, long-horizon tasks.

    Source: §6 · p. 14

    The sampling-budget analysis assumes a fixed maximum trajectory length and measures each query’s budget by tree width.

    Source: §2.2 · p. 4

Abstract

from arXiv · show

Recent advancements in aligning large language models via reinforcement learning have achieved remarkable gains in solving complex reasoning problems, but at the cost of expensive on-policy rollouts and limited exploration of diverse reasoning paths. In this work, we introduce TreePO, involving a self-guided rollout algorithm that views sequence generation as a tree-structured searching process. Composed of dynamic tree sampling policy and fixed-length segment decoding, TreePO leverages local uncertainty to warrant additional branches. By amortizing computation across common prefixes and pruning low-value paths early, TreePO essentially reduces the per-update compute burden while preserving or enhancing exploration diversity. Key contributions include: (1) a segment-wise sampling algorithm that alleviates the KV cache burden through contiguous segments and spawns new branches along with an early-stop mechanism; (2) a tree-based segment-level advantage estimation that considers both global and local proximal policy optimization. and (3) analysis on the effectiveness of probability and quality-driven dynamic divergence and fallback strategy. We empirically validate the performance gain of TreePO on a set reasoning benchmarks and the efficiency saving of GPU hours from 22\% up to 43\% of the sampling design for the trained models, meanwhile showing up to 40\% reduction at trajectory-level and 35\% at token-level sampling compute for the existing models. While offering a free lunch of inference efficiency, TreePO reveals a practical path toward scaling RL-based post-training with fewer samples and less compute. Home page locates at https://m-a-p.ai/TreePO.

1 Introduction

TreePO addresses the cost of exploring diverse reasoning paths and the difficulty of assigning sparse, delayed rewards to contributing tokens. It replaces independent rollouts with a tree-based RL pipeline for efficient sampling and more granular credit assignment.

Introduction

  • LLMs may generate thousands of tokens before receiving a sparse, delayed reward, making it difficult to explore promising paths affordably and attribute outcomes to contributing tokens.Independent trajectories also use separate KV caches and can continue exploring paths already known to be incorrect without early termination.Source: §1 · p. 2
  • Tree-based sampling averagely reduces 40% of trajectory-level inference time for baselines without sacrificing performance.The tree structure reuses shared prefixes through KV caching.Source: §1 · p. 2
  • Unlike TreeRL [11] and SPO, TreePO estimates relative advantage from coherent subtrees’ collective descendant outcomes and supports base-model training without SFT.The cited approaches primarily assign credit using parent–child value differences.Source: §1 · pp. 2–3
  • TreePO uses dynamic divergence and probability-based fallback to allocate generation budget toward diverse, promising reasoning paths.Its tree-based rollout replaces standard independent sampling and makes exploration more controllable.Source: §1 · p. 2

2 TreePO: A Tree-based Training Scheme for Policy Optimization

TreePO uses shared reasoning prefixes to motivate tree-structured sampling that avoids recomputing common work and supports targeted exploration. Its rollout and advantage methods operate on segments and shared-prefix groups.

2.1 Case Study: The Aligned Model Produces Shared Prefix

  • Across 16 stochastic rollouts per prompt at temperature 0.8, early and intermediate reasoning segments often overlap before trajectories diverge.The shared prefixes include problem interpretation, variable assignment, and initial deductions; modeling them once can avoid duplicated computation and KV-cache storage.Source: §2.1 · p. 3, Figure 2 · p. 3

2.2 Tree-based Rollout Algorithm

  • TreePO replaces independent full-trajectory decoding with queued, fixed-length segment decoding, branching prompts and continuing segments until completion, failure, or an early-stop condition.Stopped paths can seed new branches to meet the required trajectory count; the sampling budget is tree width when maximum trajectory length d × l is fixed.Source: §2.2 · p. 4
  • TreePO distributes branching budget across active paths and invokes depth-first fallback only when no path remains active and the tree has fewer than w trajectories.This is designed to balance inference utilization while preserving the opportunity for long reasoning paths.Source: §2.2 · pp. 4–5
  • Heuristic sampling prunes segments with repetitive substrings and restricts fallback candidates to stopped paths containing a formatted answer or ending in [EOS].Log probabilities from token decoding can also guide sampling without additional inference cost.Source: §2.2 · p. 5

2.3 Tree-based Advantage Estimation for Policy Optimization

  • TreePO starts from GRPO [2] with DAPO [3] modifications including clip-higher gradients, dynamic sampling, and token-level loss.Source: §2.3 · p. 5
  • TreePO estimates advantages within groups of trajectories sharing predecessor nodes, averages relative subgroup advantages, applies global variance normalization, and rejects all-correct or all-wrong queries.The subgroup method is intended to expose segment-level differences and reduce ambiguity when responses are homogeneous.Source: §2.3 · pp. 5–6

3 Experiment

The experiments report that TreePO sampling improves GRPO’s overall accuracy and training stability while reducing computation time; the tree-based advantage estimator adds further gains. The evaluation also compares tree-based and sequential sampling across reasoning benchmarks.

3.1 Hyper Parameters

  • Training experiments use Qwen2.5-7B, with Qwen2.5-7B-Instruct and Qwen2.5-Math-7B-Instruct also used to compare sampling efficiency.Training data comes from MATH and DeepScaler; evaluation uses AIME 2024, AMC 2023, MATH500, MINERVA, and Olympiad Bench, with majority-voting accuracy from 1,000 samples as the main metric.Source: §3.1 · p. 6
  • Online training tests tree depth and segment budgets of 28×256, 14×512, and 7×1024, with maximum width 16 and a branching budget of 2d per depth.The study also distinguishes More Init Divergence, which permits 2 to 8 additional initial divergences, from Fixed Init Divergence.Source: §3.1 · p. 7

3.2 Main Results

  • GRPO with TreePO sampling raises overall accuracy from 46.63% to 54.61% across datasets.The validation curves also show greater training stability than GRPO alone.Source: §3.2 · p. 7, Table 1 · p. 7
  • TreePO sampling cuts GPU hours by 12% to 43%, although it does not always beat a strong sequential baseline in final accuracy.For More Init Divergence, Sequential reaches 58.21% versus 58.06% for TreePO b=8.Source: §3.2 · p. 7
  • The TreePO advantage estimator adds 3.6% with More Init Divergence and 2.27% with Fixed Init Divergence.With tree sampling, its validation trajectory is the most stable and consistently high-performing.Source: §3.2 · p. 7

4 Discussion

The Discussion examines five research questions spanning sampling efficiency, advantage design, tree-sampling hyperparameters, efficacy–compute trade-offs, and segment-level branching control.

4.1 Sampling Efficiency Analysis

  • Both TokenPS and TrajPS peak at intermediate depth–segment settings, with model-specific optima reflecting different prefill and decoding trade-offs.Qwen2.5-7B-Instruct peaks at depth 28, Qwen2.5-Math-7B at depth 14, and Qwen2.5-Math-7B-Instruct has TokenPS peak at 14 and TrajPS peaks at 28 and 56.Source: Figure 4 · p. 8, §4.1 · p. 9
  • Under the same batch size, rollout count, and budget B, tree-based sampling yields on average +40% TrajPS and +30% TokenPS across the three models.Source: §4.1 · p. 8
  • Rollout scaling depends on model and workload: instruction-tuned models gain near-linear or stable ≈2× throughput, while Qwen2.5-Math-7B peaks around 16 rollouts before throughput declines.Shared-prefix reuse and parallel decoding support gains, while divergence, cache and memory overhead, and weaker batching can limit them.Source: Figure 5 · p. 9, §4.1 · p. 9

4.2 Analysis on the TreePO Advantage Estimation

  • Simple subgroup averaging improves accuracy and stabilizes entropy versus size-weighting, while subgroup rejection weakens learning and lengthens generations.Size-weighting can overemphasize large, easy subgroups; rejecting all-positive or all-negative subgroups removes informative calibration cases.Source: Figure 6 · p. 10, §4.2 · p. 10
  • Dropping the root-group advantage yields comparable curves, whereas a token-misaligned fallback lowers AIME accuracy and sharply increases response length.The misaligned condition uses 7 × 1024 segments with a 512-token random fallback.Source: §4.2 · pp. 10–11, Figure 6 · p. 10

4.3 Analysis of Segment Budget

  • The 14 × 512 setting achieves the highest final MATH/AIME accuracy, while 7 × 1024 lags, especially on AIME; the better accuracy comes with longer generations.The 14 × 512 setting also yields higher entropy, whereas 7 × 1024 keeps outputs shorter but sacrifices accuracy.Source: Figure 7 · p. 11, §4.3 · p. 11

4.4 Analysis of Probability-based Branching Assignment

  • Static low-probability branching underperforms the baseline and scheduled variant, while high-probability branching risks overly greedy search; greater entropy alone does not improve task performance.Low Prob Encourage has the lowest accuracy on both benchmarks and longer responses; High Prob Encourage yields the lowest entropy and shortest responses.Source: Figure 8 · p. 12, §4.4 · p. 12

4.5 Compute Scaling for Tree Sampling

  • The compute-optimal divergence strategy changes with budget: d = 2 is more efficient at low budgets, while d = 8 reaches the highest peak performance as compute increases.All tested configurations improve with compute before diminishing returns; TreePO offers distinct performance–compute curves for different divergence factors.Source: Figure 9 · p. 13, §4.5 · p. 13

5 Related Work

Related work addresses efficient rollout sampling and tree-based reasoning exploration. Prior methods reduce sampling overhead or increase branching, but the cited approaches leave limitations in trajectory awareness, credit assignment, depth control, or runtime.

Efficient Sampling

  • Batching completions and reusing prompt KV caches keep GPUs busy,,, while micro-batching and early back-propagation introduce trade-offs.Micro-batching does not inspect partial trajectories or change the advantage estimator; cutting chains early loses long-range information and complicates credit assignment.Source: §5 · p. 13

Segment-level Modeling

  • Tree-search systems build explicit rollout trees to explore multiple reasoning branches and provide denser feedback than plain chain sampling [11].The passage identifies,, [11] and as examples.Source: §5 · p. 13
  • Computing log-probabilities before branching again doubles the running time [11]; adopts an unconstrained tree and a “progress advantage” similar to Monte-Carlo returns, without depth control or validation against a frozen base policy.Source: §5 · p. 14

6 Conclusion

TreePO reformulates on-policy rollouts as segment-based tree search and uses a hierarchical advantage estimator to reduce reasoning compute while improving training stability and maintaining strong performance. The framework may support scaling reinforcement learning to more complex, long-horizon tasks.

6 Conclusion

  • TreePO combines segment-based tree search for on-policy rollouts with a hierarchical advantage estimator for complex-reasoning LLM training.The framework targets computational inefficiency and exploration instability.Source: §6 · p. 14
  • TreePO reduces reasoning computational costs, improves training stability, and maintains strong performance.Source: §6 · p. 14
  • TreePO’s efficiency and structural modeling open avenues for scaling reinforcement learning to long-horizon tasks such as multi-turn dialogue, tool use, and multi-agent systems.Source: §6 · p. 14

Related papers

Loading 2508.17445v1…