TreePO: Enhancing Policy Efficacy and Inference Efficiency with Tree Modeling

Yizhi Li, Qingshui Gu, Zhoufutu Wen, Ziniu Li, Ruibin Yuan, Tianshun Xing, Shuyue Guo, Tuney Zheng, Xin Zhou, Xingwei Qu, Wangchunshu Zhou, Zheng Zhang, Wei Shen, Wei Xue, Qian Liu, Chenghua Lin, Jian Yang, Ge Zhang, Wenhao Huang
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:66963-66976, 2026.

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-li26au, title = {{T}ree{PO}: Enhancing Policy Efficacy and Inference Efficiency with Tree Modeling}, author = {Li, Yizhi and Gu, Qingshui and Wen, Zhoufutu and Li, Ziniu and Yuan, Ruibin and Xing, Tianshun and Guo, Shuyue and Zheng, Tuney and Zhou, Xin and Qu, Xingwei and Zhou, Wangchunshu and Zhang, Zheng and Shen, Wei and Xue, Wei and Liu, Qian and Lin, Chenghua and Yang, Jian and Zhang, Ge and Huang, Wenhao}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {66963--66976}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/li26au/li26au.pdf}, url = {https://proceedings.mlr.press/v306/li26au.html}, 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.} }
Endnote
%0 Conference Paper %T TreePO: Enhancing Policy Efficacy and Inference Efficiency with Tree Modeling %A Yizhi Li %A Qingshui Gu %A Zhoufutu Wen %A Ziniu Li %A Ruibin Yuan %A Tianshun Xing %A Shuyue Guo %A Tuney Zheng %A Xin Zhou %A Xingwei Qu %A Wangchunshu Zhou %A Zheng Zhang %A Wei Shen %A Wei Xue %A Qian Liu %A Chenghua Lin %A Jian Yang %A Ge Zhang %A Wenhao Huang %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-li26au %I PMLR %P 66963--66976 %U https://proceedings.mlr.press/v306/li26au.html %V 306 %X 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.
APA
Li, Y., Gu, Q., Wen, Z., Li, Z., Yuan, R., Xing, T., Guo, S., Zheng, T., Zhou, X., Qu, X., Zhou, W., Zhang, Z., Shen, W., Xue, W., Liu, Q., Lin, C., Yang, J., Zhang, G. & Huang, W.. (2026). TreePO: Enhancing Policy Efficacy and Inference Efficiency with Tree Modeling. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:66963-66976 Available from https://proceedings.mlr.press/v306/li26au.html.

Related Material