TreePO: Enhancing Policy Efficacy and Inference Efficiency with Tree Modeling

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.

Problem: Teaching LLM to solve complex problems requires it to generate long, step-by-step reasoning paths. Current training methods are highly inefficient because they generate independent paths from scratch, repeatedly calculating the same early steps and wasting time on obvious mistakes. Solution: We introduce TreePO, which structures the LLM's reasoning like a growing tree. Instead of starting over for every path, the model shares common initial steps and branches out only at key decisions. We also use a new scoring system to precisely identify which logical branch led to the correct answer. Impact: By eliminating redundant calculations and stopping bad attempts early, TreePO saves up to 43% in computing time while keeping problem-solving accuracy high. This offers a significantly faster and greener way to develop advanced reasoning models.