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
2025-08-27
Summary
This paper introduces a new method called TreePO to improve how large language models are trained using reinforcement learning, specifically for tasks that require complex reasoning.
What's the problem?
Training these powerful language models to think through problems step-by-step is really expensive and time-consuming. Current methods require a lot of computational power to test different possible solutions, and they often don't explore enough diverse approaches to find the best answer. Essentially, it's hard to get these models to 'show their work' efficiently and thoroughly.
What's the solution?
TreePO tackles this by treating the process of generating a response as exploring a tree of possibilities. Instead of trying everything all at once, it strategically branches out and explores different paths, focusing on areas where the model is uncertain. It does this by breaking down the generation into smaller segments and only expanding on the most promising ones, saving a lot of computing power. It also uses a clever way to estimate how good each path is, considering both the overall goal and the immediate steps.
Why it matters?
This research is important because it makes it more practical to train these advanced language models. By significantly reducing the amount of computing time and resources needed, TreePO opens the door to scaling up reinforcement learning for language models, allowing them to become even better at complex reasoning tasks without breaking the bank. It also offers a way to improve existing models without retraining them from scratch.
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.