突破PPO训练效率瓶颈!字节跳动提出T-PPO,推理LLM训练速度提升2.5倍

发布于:2025-06-22 ⋅ 阅读:(23) ⋅ 点赞:(0)

突破PPO训练效率瓶颈!字节跳动提出T-PPO,推理LLM训练速度提升2.5倍

在大语言模型(LLM)通过长思维链(CoT)展现出强大推理能力的当下,强化学习(RL)作为关键技术却面临训练效率难题。本文将介绍字节跳动提出的Truncated Proximal Policy Optimization(T-PPO),其通过创新方法提升训练效率,在AIME 2024基准测试中表现亮眼,一起来了解这一重要进展。

论文标题
Truncated Proximal Policy Optimization

来源
arXiv:2506.15050v1 [cs.AI] + https://arxiv.org/abs/2506.15050

文章核心

研究背景

近年来,推理导向的大型语言模型(LLM)如OpenAI的o1、DeepSeekR1和QwQ等,借助扩展的思维链(CoT)推理在数学推理、编程和基于代理的任务等复杂领域展现出最先进的性能,而深度强化学习(RL)技术是这些模型提升推理能力的重要支撑。

研究问题

  1. PPO的在线策略特性导致训练效率低下:PPO作为LLM优化的主要RL方法,其在线策略本质限制了训练效率,在处理长CoT轨迹时,这种限制尤为明显,会导致大量计算开销和延长训练时间。

  2. 长生成过程中硬件利用率低:完全同步的长生成过程中,资源在等待完整rollout期间经常处于闲置状态,存在硬件利用率低的固有缺点。

  3. 离线策略方法存在训练不稳定问题:虽然离线策略方法训练效率更高,但通常在策略梯度估计器中存在高方差,导致训练不稳定和性能下降。

主要贡献

  1. 提出EGAE进行优势估计:开发Extended Generalized Advantage Estimation(EGAE),可从不完整响应中进行优势估计,同时保持策略学习的完整性,使策略更新能在轨迹完全生成前进行,提高计算资源利用率。

  2. 设计计算优化机制:创建一种计算优化机制,允许策略和价值模型独立优化,通过选择性过滤提示和截断令牌,减少冗余计算并加速训练过程,且不牺牲收敛性能。

  3. 提升训练效率与性能:在AIME 2024上使用32B基础模型的实验表明,T-PPO将推理LLM的训练效率提高了2.5倍,性能优于现有竞争对手,在AIME’24基准测试中取得62的pass@1分数。

方法论精要

核心算法/框架

T-PPO是PPO的新型扩展,核心在于EGAE和令牌过滤策略。EGAE扩展了传统的GAE,支持使用部分生成的响应进行策略优化;令牌过滤策略通过截断生成和选择性使用令牌,实现策略和价值模型的独立优化。

(所以EGAE的关键是,对于未生成的tokens V ( s l ) = V ( s l − 1 ) V(s_{l})=V(s_{l-1}) V(sl)=V(sl1),算是一种近似。那 δ t \delta _ t δt里的 r t r_t rt是怎么来的?kl?)

关键参数设计原理

窗口长度 l l l:用于截断生成,假设实际最大响应长度 L L L与窗口长度 l l l的比值为 k k k,生成时间和训练时间大约可节省k倍。

EGAE中的参数 λ λ λ γ γ γ λ λ λ控制未来奖励对优势估计的影响, γ γ γ为折扣因子,通过调整它们控制偏差-方差权衡。

裁剪参数:策略的 ϵ l o w = 0.2 \epsilon_{low}=0.2 ϵlow=0.2 ϵ h i g h = 0.28 \epsilon_{high}=0.28 ϵhigh=0.28,价值函数的 ξ l o w = 0.5 \xi_{low}=0.5 ξlow=0.5 ξ h i g h = 0.6 \xi_{high}=0.6 ξhigh=0.6,限制更新幅度以保证稳定性。

创新性技术组合

将EGAE与令牌过滤策略结合,实现不完整轨迹的优势计算和渐进式策略更新。

策略模型训练使用当前训练步骤生成的响应令牌,价值模型训练使用完成序列的所有生成令牌,且价值模型采用蒙特卡洛训练范式以确保无偏估计。

采用连续批处理策略,当某些序列达到结束条件时,在下一步插入新提示,未完成样本保留,保持每步批大小恒定。

实验验证方式

数据集:使用美国数学邀请赛(AIME)作为推理问题的代表性基准,训练集为DAPO-Math-17K,包含过去所有AIME竞赛问题及一些人工构造的难题。

基线方法:对比DeepSeek-R1-Zero-Qwen-32B、DAPO、VAPO、GePPO、PPO-EWMA等,通过AIME 2024的pass@1分数和训练时间评估性能。

实验洞察

性能优势

T-PPO在AIME 24上实现61.88的pass@1分数,超越DeepSeek-R1-Zero-Qwen-32B和现有的最佳异步PPO算法,在相同性能下,与需要20k响应长度的PPO相比,在AIME24基准上wall-clock time减少60%。

效率突破

T-PPO的平均每1000步壁钟时间与PPO-EWMA相当,远低于vanillaPPO算法,且收敛步骤(6720步)显著少于PPO-EWMA(11200步),总运行时间更短;在policy rollout中的计算强度为249 operations/byte,远高于PPO的84 operations/byte,更好地利用了计算资源。

训练动态分析

对响应长度的分析表明,其呈现先增加、暂时下降、再恢复并最终稳定的特征,最终稳定的响应长度超过vanillaPPO,说明T-PPO保留并可能增强了推理模型的长度缩放能力,模型在学习过程中不断完善推理方法。


网站公告

今日签到

点亮在社区的每一天
去签到