arXiv:2505.15201cs.LGcs.AI2025-05NeurIPS被引 76

让强化学习同时优化多个尝试的综合表现,突破传统方法瓶颈。

Pass@K Policy Optimization: Solving Harder Reinforcement Learning Problems

  • 设计新奖励机制,直接优化 pass@k 性能
  • 支持任意 k ≤ n,且训练中可动态调整 k 值
  • 在复杂任务上显著提升解题能力,适合难样本学习

强化学习算法通常对每个问题采样多个解法(n>1),并独立奖励,这仅优化 pass@1 性能,忽视了样本集的多样性与协同价值,限制了探索能力。为此,本文提出 Pass-at-k Policy Optimization(PKPO),通过变换最终奖励,直接优化 pass@k 表现,使模型关注一组样本的联合收益。我们推导出 pass@k 及其梯度的低方差无偏估计器,适用于二元与连续奖励场景。优化过程转化为标准 RL,只需对奖励进行稳定高效的函数变换。相比以往仅支持 k=n,本方法首次实现任意 k ≤ n 的稳健 pass@k 优化。此外,可通过训练中渐进式调整 k 值,在不牺牲 pass@1 的前提下显著提升 pass@k。实验验证了该方法的降方差特性;在 GEMMA-2 等真实任务中,更高 k 值可解决更多更难问题,且 annealing k 能同时提升 pass@1 与 pass@k。对于传统 pass@1 方法陷入停滞的挑战性任务集,pass@k 方法有效激活学习,源于其优先考虑联合效用而非单个样本性能。

原文摘要 · Abstract (English)

Reinforcement Learning (RL) algorithms sample multiple n>1 solution attempts for each problem and reward them independently. This optimizes for pass@1 performance and prioritizes the strength of isolated samples at the expense of the diversity and collective utility of sets of samples. This under-utilizes the sampling capacity, limiting exploration and eventual improvement on harder examples. As a fix, we propose Pass-at-k Policy Optimization (PKPO), a transformation on the final rewards which leads to direct optimization of pass@k performance, thus optimizing for sets of samples that maximize reward when considered jointly. Our contribution is to derive novel low variance unbiased estimators for pass@k and its gradient, in both the binary and continuous reward settings. We show optimization with our estimators reduces to standard RL with rewards that have been jointly transformed by a stable and efficient transformation function. While previous efforts are restricted to k=n, ours is the first to enable robust optimization of pass@k for any arbitrary k <= n. Moreover, instead of trading off pass@1 performance for pass@k gains, our method allows annealing k during training, optimizing both metrics and often achieving strong pass@1 numbers alongside significant pass@k gains. We validate our reward transformations on toy experiments, which reveal the variance reducing properties of our formulations. We also include real-world examples using the open-source LLM, GEMMA-2. We find that our transformation effectively optimizes for the target k. Furthermore, higher k values enable solving more and harder problems, while annealing k boosts both the pass@1 and pass@k . Crucially, for challenging task sets where conventional pass@1 optimization stalls, our pass@k approach unblocks learning, likely due to better exploration by prioritizing joint utility over the utility of individual samples.

强化学习奖励设计多样本优化LLM推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。