Paper 2508.17445
TreePO: Bridging the Gap of Policy Optimization and Efficacy and Inference Efficiency with Heuristic Tree-based Modeling
- Published
- Aug 2025
- Research lab
- Independent
- Citations
- 59
- GitHub
- 65 stars
01 In brief
Summary
TreePO is a reinforcement learning framework for large language models that replaces standard independent trajectory sampling with a heuristic tree-based rollout mechanism.
It models sequence generation as a tree search with dynamic branching and fixed-length segment decoding, leveraging shared prefixes via KV-cache reuse to reduce compute.
The method introduces a tree-based segment-level advantage estimation that aggregates subgroup rewards for more precise credit assignment, and it supports training from base models without supervised fine-tuning.
Experiments on math reasoning benchmarks (AIME, AMC, MATH, MINERVA, Olympiad Bench) show that TreePO improves overall accuracy from 46.63% to 54.61% over GRPO, and when combined with the advantage estimator, reaches 58.21% with more initial divergence.
It also reduces GPU hours by 12% to 43% and trajectory-level inference time by up to 40%.
Ablations reveal that simple averaging of subgroup advantages outperforms size-weighting, subgroup-level rejection hurts performance, and token-aligned segments are crucial.
Probability-based branching control does not improve performance, and compute scaling shows that optimal tree divergence depends on the compute budget.
TreePO offers a more efficient and stable alternative for RL-based post-training of reasoning models.
02 From the paper
Abstract
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.