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

**Authors:** 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.17445](https://arxiv.org/abs/2508.17445)  
**Date:** 2025-08-24  
**Summary page:** https://www.summarizepaper.com/en/paper/2508.17445v1/

## Key points

- Reinforcement Learning (RL) enhances reasoning abilities of Large Language Models (LLMs)
- Challenges in RL for LLMs: exploration and exploitation, long sequences before feedback, diverse exploration paths while reducing computational costs and accurately attributing rewards
- Traditional RL approaches are inefficient and limit adaptability
- Monte Carlo Tree Search (MCTS) is promising but can be inefficient for LLMs due to sequential rollouts
- Tree-based Policy Optimization (TreePO) framework leverages tree structures to improve sampling efficiency and credit assignment in RL pipelines
- TreePO replaces independent rollouts with a self-guided tree search, reducing trajectory-level inference time by 40%
- Key contributions of TreePO: segment-wise sampling algorithm, tree-based segment-level advantage estimation, probability-driven dynamic divergence strategies
- Empirical validation shows performance gains on reasoning benchmarks and efficiency savings up to 43% in GPU hours
- TreePO offers inference efficiency without sacrificing performance, enabling scaling of RL-based post-training with fewer samples and less compute

## Summary

, , , , 

Reinforcement Learning (RL) has shown great promise in enhancing the reasoning abilities of Large Language Models (LLMs), but faces challenges related to exploration and exploitation. LLMs must generate long sequences before receiving feedback, leading to questions about how to enable diverse exploration paths while reducing computational costs and accurately attributing rewards. Traditional RL approaches generate multiple independent trajectories for a single query, which is inefficient and limits adaptability. Monte Carlo Tree Search (MCTS) offers a promising alternative, but can be inefficient for LLMs due to sequential rollouts. In response to these challenges, we introduce Tree-based Policy Optimization (TreePO), a framework that leverages tree structures to improve sampling efficiency and credit assignment in RL pipelines. TreePO replaces independent rollouts with a self-guided tree search that maximizes reuse of shared prefixes, reducing trajectory-level inference time by an average of 40%. Our approach facilitates granular advantage estimation, enabling robust credit assignment based on collective outcomes of sub-trees. Key contributions of TreePO include: a segment-wise sampling algorithm that reduces the KV cache burden through contiguous segments and early-stop mechanisms; a tree-based segment-level advantage estimation that considers global and local policy optimization; and analysis on probability-driven dynamic divergence strategies. Empirical validation demonstrates performance gains on reasoning benchmarks and efficiency savings of up to 43% in GPU hours for trained models. By offering inference efficiency without sacrificing performance, TreePO paves the way for scaling RL-based post-training with fewer samples and less compute. Overall, our work addresses fundamental challenges in RL for LLMs by introducing a structured sampling mechanism that improves computational efficiency and credit assignment accuracy. Through TreePO, we demonstrate the potential for more effective reasoning capabilities without the need for supervised fine-tuning, positioning our framework as a practical solution for advancing complex reasoning tasks in language models.