TreePO: Bridging the Gap of Policy Optimization and Efficacy and
  Inference Efficiency with Heuristic Tree-based Modeling
Authors: name: Yizhi Li; name: Qingshui Gu; name: Zhoufutu Wen; name: Ziniu Li; name: Tianshun Xing; name: Shuyue Guo; name: Tianyu Zheng; name: Xin Zhou; name: Xingwei Qu; name: Wangchunshu Zhou; name: Zheng Zhang; name: Wei Shen; name: Qian Liu; name: Chenghua Lin; name: Jian Yang; name: Ge Zhang; name: Wenhao Huang
Published: 2025-08-24 16:52:37

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.