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
TL;DR
RL training for complex reasoning is constrained by expensive on-policy rollouts, limited exploration, sparse delayed rewards, and difficult credit assignment. TreePO reformulates rollouts as a heuristic, segment-based tree search with shared-prefix reuse, dynamic branching, fallback strategies, and tree-based advantage estimation. The paper reports reduced reasoning computation, improved training stability, and strong performance, supporting a more efficient and scalable training frontier.
Problem
Long-horizon RL reasoning requires diverse exploration and accurate credit assignment despite expensive rollouts and sparse, delayed rewards.
Method
TreePO replaces independent rollouts with a heuristic tree-based search that reuses shared prefixes, dynamically branches or falls back, and estimates advantages over sub-trees.
Results
TreePO reduces reasoning computational costs while improving training stability and maintaining strong performance, yielding a superior computational-cost and model-performance trade-off.
Takeaways & Limitations
TreePO offers a practical route toward scaling RL post-training to more complex, long-horizon tasks with fewer samples and less compute.
Takeaways & Limitations
The sampling-budget analysis assumes fixed maximum trajectory length and characterizes the budget through tree width; an additional displayed constraint limits the supported setting.
Abstract
from arXiv · showhide
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 exploration, computational-efficiency, and credit-assignment challenges of long-horizon RL reasoning by replacing independent rollouts with a heuristic tree-based process. Its tree sampling and advantage estimation improve exploration and computational efficiency while supporting scalable training.
- Motivation: Long reasoning sequences create sparse, delayed rewards, making diverse path exploration and token-level credit assignment difficult.The paper frames these as two central research challenges: exploring potentially correct paths without increasing costs and attributing outcome rewards to contributing tokens.
- Motivation: Independent trajectories waste computation through separate KV caches and continue exploring paths that may already be incorrect.The paper identifies both inefficient cache use and the lack of early termination as limitations of standard sampling.
- Approach: TreePO structures rollouts as a heuristic tree search that reuses shared prefixes, supports early termination, and allocates exploration through dynamic divergence and probability-based fallback.The design targets both inference efficiency and controllable exploration by replacing independent rollouts with a unified tree-based process.
- Approach: TreePO estimates advantages over coherent sub-trees using collective descendant outcomes rather than only parent-child value differences.This provides a more granular relative credit-assignment mechanism and is designed to train directly from a base model without initial instruction tuning.
- Results: TreePO provides a superior computational-cost and model-performance trade-off, establishing a more efficient and scalable training frontier.The introduction reports this conclusion from extensive experiments across reasoning benchmarks.
2 TreePO: A Tree-based Training Scheme for Policy Optimization
TreePO reorganizes reasoning rollouts as a heuristic, segment-level search tree that reuses shared prefixes, prunes failed paths, and supports finer-grained policy optimization.
- Motivation: 16 stochastic rollouts reveal extensive shared reasoning segments, motivating tree-structured sampling instead of recomputing common prefixes independently.Shared prefixes create duplicated computation and KV-cache storage in standard on-policy reinforcement learning.
- Segment-level Tree Sampling: TreePO decodes fixed-length segments, maintains a prompt queue, and stops branches at end-of-sequence, segment-length, or designed early-stopping criteria.Completed or failed paths are removed from the active queue while unfinished paths receive concatenated segments as new prompts.
- Branching and Fallback: TreePO uses branching and fallback policies to balance rollout efficiency, exploration, and the requirement to obtain the target tree width.Fallback begins when no active path remains and the tree still has fewer than the required number of trajectories.
- Heuristic Sampling: Heuristic controls dynamically allocate search effort through early stopping, branching, and fallback, including pruning repetitive or mumbling paths.The method also transfers branching budget across active paths to avoid low-utilization batches and skew toward short paths.
- Policy Optimization: TreePO adopts GRPO with DAPO modifications as its optimization starting point, including clipped policy ratios, dynamic sampling, and token-level loss.The formulation also requires sampled outputs to be equivalent to the reference answer under the stated constraint.
- Tree-based Advantage Estimation: TreePO divides trajectories into inference-step segments and estimates advantages using trajectory subgroups that share predecessor nodes.This segment-level structure captures intra-response differences that global advantages can obscure when responses are similar.
3 Experiment
Experiments evaluate TreePO across reasoning benchmarks, training configurations, and sampling settings. Tree-based sampling improves stability and efficiency, while the advantage estimator adds further gains, with some trade-off in peak accuracy or convergence speed.
- Tree Setting: The study searches tree depth and segment budgets across {28×256, 14×512, 7×1024}, with binary branching and maximum width 16.The listed settings vary tree depth and segment token budget under a fixed response-length constraint.
- Main Results: Tree-based methods enhance training stability and computational efficiency, although some configurations trade off raw convergence speed and peak accuracy.This conclusion is based on the main results and training curves.
- Main Results: 46.63% to 54.61%: adding TreePO sampling to GRPO increases overall accuracy across datasets and produces more stable validation behavior.The comparison is reported for GRPO with TreePO Sampling versus GRPO.
- Main Results: 12% to 43%: tree-based sampling reduces GPU hours, although it does not always exceed a strong sequential baseline in final accuracy.One reported comparison is 58.21% for Sequential versus 58.06% for TreePO b=8 in the More Init Divergence model.
- Main Results: 3.6% and 2.27%: the TreePO advantage estimator further improves training with More Init Divergence and Fixed Init Divergence, respectively.The estimator is associated with the most stable and consistently high-performing validation trajectory.
4 Discussion
The discussion evaluates TreePO across efficiency, advantage estimation, segment configuration, and branching controls. Results show substantial throughput gains, configuration-sensitive optimization, and trade-offs between exploration, accuracy, and response length.
- Sampling efficiency: Tree-based sampling improves throughput by averaging +40% TrajPS and +30% TokenPS across three models under matched rollout settings.The efficiency comparison uses the same batch size, rollout count, and budget B.
- Sampling efficiency: Both TokenPS and TrajPS peak at intermediate depth–segment combinations because shorter segments increase recomputation while deeper trees improve branching parallelism.Prefill favors longer segments and shallower trees, whereas decoding benefits from deeper trees and more parallel rollouts.
- Sampling efficiency: Rollout scaling is model- and workload-dependent: some models approach 2× throughput, while divergent trajectories eventually reduce cache reuse and batching efficiency.Qwen2.5-7B-Instruct reaches roughly 2×, Qwen2.5-Math-7B-Instruct maintains about 2×, and Qwen2.5-Math-7B peaks around 16 rollouts before declining.
- Advantage estimation: Simple subgroup averaging yields higher MATH and AIME accuracy with lower, more stable entropy and no unnecessary response-length growth than size-weighted aggregation.Size-weighting over-emphasizes large, easy subgroups and down-weights informative small, hard subgroups.
- Advantage estimation: Removing the root-group advantage produces comparable curves, suggesting aggregated subgroup signals can approximate the full-group optimization signal.The result indicates the root term is not strictly necessary in the tested setting.
- Segment budget and alignment: Token-misaligned fallback sharply lowers AIME accuracy and increases response length, while 14 × 512 outperforms 7 × 1024 under group-size-weighted advantage.The findings support token-aligned segments and deeper trees with moderate segment lengths for stable optimization and credit assignment.
- Segment budget and alignment: Higher accuracy is coupled with longer generations, while shorter outputs from 7 × 1024 sacrifice accuracy.The best-performing 14 × 512 configuration also produces the largest response-length growth and higher entropy.
- Branching control: Static low- or high-probability branching controls underperform the baseline, and forcing low-probability exploration increases entropy without improving benchmark performance.The low-probability strategy is associated with the lowest accuracy, longer responses, and less effective reasoning paths.
5 Related Work
Related work develops efficient batched sampling and tree- or segment-based reinforcement learning, but existing approaches retain important inference or modeling limitations that motivate TreePO.
- Efficient Sampling: Efficient-sampling methods batch completions, reuse prompt KV caches, and parallelize token groups to reduce rollout-loop overhead.Some methods further divide groups into micro-groups and continuously interleave execution.
- Segment-level Modeling: Tree-search reinforcement-learning systems explore multiple reasoning branches and provide denser feedback than plain chain sampling.TreeRL, SPO, and related methods use explicit tree or segment structures for exploration and credit assignment.
- Segment-level Modeling: Prior approaches can suffer from shallow trees, full-answer recomputation before branching, unconstrained depth, or limited validation against frozen base policies.These limitations include doubled running time in one system and insufficient depth control in another.
6 Conclusion
TreePO reformulates on-policy reinforcement-learning rollouts as segment-based tree search with hierarchical advantage estimation. The framework is presented as reducing reasoning cost, improving training stability, and maintaining strong performance while supporting longer-horizon applications.
- Conclusion: TreePO addresses computational inefficiency and exploration instability by combining segment-based tree search with hierarchical advantage estimation.The framework targets reinforcement-learning training for large language models performing complex reasoning.
- Conclusion: The authors report reduced reasoning computational costs, improved training stability, and strong performance from TreePO.They identify multi-turn dialogue, tool use, and multi-agent systems as promising longer-horizon application areas.
Authorship
The paper lists authors affiliated with M-A-P, ByteDance Seed, and the University of Manchester.
- Affiliations: The author list includes contributors from M-A-P, ByteDance Seed, and the University of Manchester.The affiliations are labeled 1, 2, and 3 respectively.