Pass@K 策略优化：求解更难的强化学习问题

强化学习（RL）算法为每个问题采样多个 $n>1$ 次尝试并独立给予奖励。这优化了 pass@1 性能，优先考虑孤立样本的强度，而牺牲了样本集的多样性和集体效用。这未能充分利用采样能力，限制了探索和最终在更难样本上的提升。作为解决方案，我们提出了 Pass-at-k 策略优化（PKPO），一种对最终奖励的变换，可直接优化 pass@k 性能，从而优化在联合考虑时最大化奖励的样本集。我们的贡献在于推导了 pass@k 及其梯度的新型低方差无偏估计器，适用于二元奖励和连续奖励两种设置。我们证明，使用我们的估计器进行优化可归结为标准 RL，其奖励经过稳定高效的变换函数联合变换。虽然先前工作仅限于 $k=n$， ours 是首个支持对任意 $k \leq n$ 进行稳健 pass@k 优化的方法。此外，我们的方法不是以 pass@1 性能换取 pass@k 增益，而是允许在训练期间对 $k$ 进行退火，同时优化两个指标，通常还能取得较强的 pass@1 数值以及显著的 pass@k 提升。我们在玩具实验上验证了奖励变换，揭示了所提公式的方差降低特性。我们还包括使用开源大语言模型 GEMMA-2 的真实世界示例。我们发现变换有效优化了目标 $k$。此外，更高的 $k$ 值能够求解更多更难的问题，而 $k$ 的退火同时提升了 pass@1 和 pass@k。关键的是，对于传统 pass@1 优化停滞的具有挑战性的任务集，我们的 pass@k 方法解除了学习瓶颈，可能是因为通过优先考虑样本集的联合效用而非单个样本的效用，实现了更好的探索。