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

TL;DR

TreePO reframes policy optimization for reasoning LLMs as a heuristic tree-based rollout, enabling shared-prefix KV-cache reuse and early pruning to cut compute while maintaining exploration. It introduces segment-level tree sampling, dynamic branching with fallback, and a tree-aware advantage estimation that aggregates rewards across subtrees. Empirical results show substantial reductions in GPU hours (12-43%) and improvements in stability and accuracy on reasoning benchmarks, with model- and workload-dependent sweet spots for depth and segment configurations. The approach offers a practical path toward scalable RL-based post-training with reduced sample complexity and compute, and situates itself among efficient sampling and segment-level RL methods. Overall, TreePO demonstrates how structured rollout and hierarchical credit assignment can unlock faster, more diverse reasoning in large language models.

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.

Figures (9)

• Figure 1: Demonstration of the Validation Performance Curves along Training based on Qwen2.5-7B (Left, Mid) and Demonstration of TreePO Sampling (Right). Left, Mid: Compared to the GRPO setting, although replaced additional treed-based sampling causes a slower convergence, it could stabilize the training. When cooperate the health.
• Figure 2: Multiple sampled trajectories from the same prompt, with shared reasoning segments highlighted in matching colors. Despite stochastic generation, key problem-solving steps are consistently reproduced.
• Figure 3: Demonstration of the TreePO Advantage Estimation. Assuming that the tree-based sampling has derived 8 trajectories (leaf nodes) given a query $q$, we take node $c_{2,2}$ as an example to calculate the sub-group advantages. The tree-based sub-groups could be further defined by its predecessors $c_{2}$, $c$, and $q$. Thus the final advantages can be calculated as the averagely aggregated sub-group advantages.
• Figure 4: Performance comparison between Tree-based Sampling and Conventional Sampling across different tree depths.
• Figure 5: Performance comparison between Tree-based Sampling and Conventional Sampling across different numbers of rollouts.