TreePO：通过启发式树结构建模桥接策略优化与有效性及推理效率之间的差距
机器学习
2025-08-26 v1 计算与语言
摘要
通过强化学习对大型语言模型进行对齐的最新进展在解决复杂推理问题方面取得了显著的成果，但以高昂的 on-policy rollout 成本和有限的多样化推理路径探索为代价。本文引入了 TreePO，涉及一种自指导的 rollout 算法，将序列生成视为树结构搜索过程。TreePO 由动态树抽样策略和固定长度分段解码组成，利用局部不确定性确保额外的分支。通过对常见前缀进行计算摊销并提前修剪低价值路径，TreePO 本质上在保持或增强探索多样性的同时大幅降低了每个更新的计算负担。关键贡献包括：(1) 一个分段抽样算法，通过连续分段和早停机制减轻 KV 缓存负担并生成新的分支；(2) 基于树的分段级优势估计，考虑全局和局部近端策略优化；(3) 关于概率和质量驱动的动态分支与回退策略的有效性分析。我们在一组推理基准测试上实证验证了 TreePO 的性能提升，显示出了从 22% 最高可达 43% 的 GPU 时间节省，同时在轨迹级和 token 级采样计算上分别实现了最高 40% 和 35% 的降低。虽然 TreePO 为推理效率带来了“礼物”，但它揭示了一条以更少样本和更少计算量实现 RL 基于后训练的可行路径。主页位于 https://m-a-p.ai/TreePO
引用
@article{arxiv.2508.17445,
title = {TreePO: Bridging the Gap of Policy Optimization and Efficacy and Inference Efficiency with Heuristic Tree-based Modeling},
author = {Yizhi Li and Qingshui Gu and Zhoufutu Wen and Ziniu Li and Tianshun Xing and Shuyue Guo and Tianyu Zheng and Xin Zhou and Xingwei Qu and Wangchunshu Zhou and Zheng Zhang and Wei Shen and Qian Liu and Chenghua Lin and Jian Yang and Ge Zhang and Wenhao Huang},
journal= {arXiv preprint arXiv:2508.17445},
year = {2025}
}