TreePO用树结构提升推理模型训练效率,省下超四成计算资源。
TreePO: Bridging the Gap of Policy Optimization and Efficacy and Inference Efficiency with Heuristic Tree-based Modeling
- 将生成过程建模为树搜索,动态扩展分支并早停低效路径。
- 在多个基准上减少22%至43%的采样耗时,轨迹级节省40%算力。
- 适合追求高效强化学习微调的研究者与工业部署场景。
近期基于强化学习对齐大语言模型在解决复杂推理问题上取得显著进展,但伴随高昂的在线策略采样成本和有限的路径探索。本文提出TreePO,通过自引导采样算法将序列生成视为树状搜索过程。其由动态树采样策略与固定长度段解码构成,利用局部不确定性触发新增分支。通过共享公共前缀计算并提前剪枝低价值路径,有效降低每次更新的计算开销,同时保持或增强探索多样性。关键贡献包括:(1) 分段采样算法通过连续片段缓解KV缓存压力,并结合早停机制生成新分支;(2) 基于树结构的分段级优势估计,兼顾全局与局部近端策略优化;(3) 对概率与质量驱动的动态发散及回退策略的有效性分析。在多个推理基准上实证验证性能提升,采样设计节省22%至43%的GPU小时;现有模型在轨迹级和词元级分别减少40%与35%的采样计算量。TreePO在不牺牲效率的前提下,为以更少样本和更低算力实现强化学习后训练提供了可行路径。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。