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
Organization: ByteDance Seed

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.