通过追踪梯度曲率,让大模型推理强化学习更稳定高效。
Stabilizing Policy Gradients for Sample-Efficient Reinforcement Learning in LLM Reasoning
- 引入曲率信息动态调整策略更新,提升优化稳定性
- 仅剔除不足8%的token,样本效率最高提升30倍
- 适合追求高效微调的大模型推理训练场景
强化学习,尤其是基于策略梯度的方法,在赋予大语言模型推理能力方面发挥了核心作用。然而,该场景下策略梯度优化的稳定性尚未得到充分研究。因此,现有方法常采用保守的超参数以保证稳定,导致需要更多训练样本并增加计算成本。本文通过显式考虑二阶几何结构,形式化了策略梯度的随机优化问题,提出一个可计算的框架,用于跟踪并利用更新过程中的曲率信息。进一步地,该框架被用于设计数据选择干预机制。所提出的算法Curvature-Aware Policy Optimization(CAPO)能够识别导致不稳定更新的样本并将其屏蔽。理论上,在合理假设下建立了单调改进保证。在标准数学推理基准上,实验表明,当基线方法崩溃时,CAPO仍能保持稳定更新。仅需极少干预(剔除少于8%的token),CAPO相比标准GRPO在大模型推理中实现最高30倍的样本效率提升。
原文摘要 · Abstract (English)
Reinforcement Learning, particularly through policy gradient methods, has played a central role in enabling reasoning capabilities of Large Language Models. However, the optimization stability of policy gradients in this setting remains understudied. As a result, existing implementations often resort to conservative hyperparameter choices to ensure stability, which requires more training samples and increases computational costs. Hence, developing models for reliably tracking the underlying optimization dynamics and leveraging them into training enables more sample-efficient regimes and further unleashes scalable post-training. We address this gap by formalizing the stochastic optimization problem of policy gradients with explicit consideration of second-order geometry. We propose a tractable computational framework that tracks and leverages curvature information during policy updates. We further employ this framework to design interventions in the optimization process through data selection. The resultant algorithm, Curvature-Aware Policy Optimization (CAPO), identifies samples that contribute to unstable updates and masks them out. Theoretically, we establish monotonic improvement guarantees under realistic assumptions. On standard math reasoning benchmarks, we empirically show that CAPO ensures stable updates under aggressive learning regimes where baselines catastrophically fail. With minimal intervention (rejecting fewer than 8% of tokens), CAPO achieves up to 30x improvement in sample efficiency over standard GRPO for LLM reasoning.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。