摘要

通过强化学习对齐大型语言模型的最新进展，在解决复杂推理问题方面取得了显著成效，但同时也带来了高昂的在线策略采样开销，以及对多样化推理路径探索不足的局限性。在本工作中，我们提出TreePO，一种采用自引导采样算法的方法，将序列生成视为一种树状结构的搜索过程。TreePO由动态树采样策略与固定长度片段解码构成，利用局部不确定性来触发新增分支。通过在共用前缀上摊销计算成本，并提前剪枝低价值路径，TreePO在保持甚至提升探索多样性的同时，显著降低了每次更新的计算负担。主要贡献包括：（1）一种分段采样算法，通过连续片段缓解KV缓存压力，并结合早期停止机制生成新分支；（2）一种基于树结构的分段级优势估计方法，兼顾全局与局部的近端策略优化；（3）对概率驱动与质量驱动的动态分叉及回退策略有效性的深入分析。我们在一组推理基准上实证验证了TreePO的性能提升，同时在训练模型的采样设计中，GPU小时数节省高达22%至43%。此外，对于现有模型，轨迹层级的采样计算减少最多达40%，token层级减少最多达35%。TreePO在不增加额外成本的前提下，显著提升了推理效率，为以更少样本和更低计算开销实现基于强化学习的后训练扩展提供了切实可行的路径。项目主页详见：https://m-a-p.ai/TreePO。