TreePO：通过启发式树形建模弥合策略优化、有效性和推理效率之间的差距

通过强化学习对大型语言模型进行对齐的最新进展在解决复杂的推理问题方面取得了显著的进步，但其代价是昂贵的 on-policy rollouts 和对不同推理路径的有限探索。在这项工作中，我们提出了 TreePO，它包含一种自导引 rollout 算法，将序列生成视为一个树状搜索过程。TreePO 由动态树采样策略和固定长度片段解码组成，利用局部不确定性来保证额外的分支。通过分摊常见前缀的计算成本并及早修剪低价值路径，TreePO 基本上降低了每次更新的计算负担，同时保持或增强了探索的多样性。主要贡献包括：(1) 一种分段采样算法，通过连续分段减轻 KV 缓存负担，并沿途生成新分支并带有早期停止机制；(2) 一种基于树的分段级别优势估计，同时考虑了全局和局部近端策略优化；(3) 对概率驱动和质量驱动动态发散及回退策略有效性的分析。我们在多个推理基准上实证验证了 TreePO 的性能提升，并在 GPU 小时效率方面节省了 22% 至 43% 的采样设计，同时对于现有模型，在轨迹级别和 token 级别的采样计算上分别降低了 40% 和 35%。TreePO 在提供推理效率的免费午餐的同时，也展示了一条通往使用更少样本和更少计算量来扩展基于 RL 的后期训练的实用途径。主页位于 https://m-a-p.ai/TreePO。