T-PPO提升大模型推理训练效率,支持长思维链生成
Truncated Proximal Policy Optimization
- 通过截断响应和不完整反馈估计优势,实现高效策略更新
- 在AIME 2024上使训练效率提升2.5倍,保持收敛性能
- 适合需要长推理链的大型语言模型训练场景
近期,测试时扩展的大规模语言模型(LLMs)通过生成长思维链(CoT)在科学与专业任务中展现出卓越的推理能力。作为构建此类推理模型的关键技术,强化学习(RL)如近端策略优化(PPO)允许模型通过试错学习。然而,由于PPO固有的在线策略特性,且响应长度不断增长,导致训练耗时严重。本文提出截断近端策略优化(T-PPO),通过简化策略更新与限制响应长度来提升训练效率。T-PPO缓解了全同步长生成过程中硬件利用率低的问题,避免资源在等待完整回放期间闲置。主要贡献有二:一是提出扩展广义优势估计(EGAE),基于不完整响应进行优势估计,同时保证策略学习完整性;二是设计计算优化机制,实现策略与价值模型独立优化。通过选择性过滤提示与截断标记,该机制减少冗余计算,加速训练过程而不牺牲收敛性能。在32B基模型上对AIME 2024的实验表明,T-PPO可将推理型LLM的训练效率提升最高达2.5倍,并优于现有竞争方法。
原文摘要 · Abstract (English)
Recently, test-time scaling Large Language Models (LLMs) have demonstrated exceptional reasoning capabilities across scientific and professional tasks by generating long chains-of-thought (CoT). As a crucial component for developing these reasoning models, reinforcement learning (RL), exemplified by Proximal Policy Optimization (PPO) and its variants, allows models to learn through trial and error. However, PPO can be time-consuming due to its inherent on-policy nature, which is further exacerbated by increasing response lengths. In this work, we propose Truncated Proximal Policy Optimization (T-PPO), a novel extension to PPO that improves training efficiency by streamlining policy update and length-restricted response generation. T-PPO mitigates the issue of low hardware utilization, an inherent drawback of fully synchronized long-generation procedures, where resources often sit idle during the waiting periods for complete rollouts. Our contributions are two-folds. First, we propose Extended Generalized Advantage Estimation (EGAE) for advantage estimation derived from incomplete responses while maintaining the integrity of policy learning. Second, we devise a computationally optimized mechanism that allows for the independent optimization of the policy and value models. By selectively filtering prompt and truncated tokens, this mechanism reduces redundant computations and accelerates the training process without sacrificing convergence performance. We demonstrate the effectiveness and efficacy of T-PPO on AIME 2024 with a 32B base model. The experimental results show that T-PPO improves the training efficiency of reasoning LLMs by up to 2.5x and outperforms its existing competitors.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。