技术摘要：Pass@K 策略优化 (PKPO)

问题陈述

应用于大语言模型（LLMs）的强化学习（RL）算法通常会对单个问题采样多个解尝试（$n > 1$），但往往基于单个样本的表现（pass@1）进行优化。这种方法优先考虑单个高奖励样本的效用，而非集合的集体效用，从而有效地降低了对采样能力的利用。因此，标准的策略梯度方法可能无法充分探索解空间，特别是在单次采样正确概率较低的困难任务上。虽然存在推理时搜索方法，但推理与模型更新之间的耦合意味着，对于那些可以验证正确性但解并非预先已知的任务，单纯优化期望单样本奖励是次优的。

方法论

作者提出了 Pass-at-k 策略优化 (PKPO) 框架，该框架直接优化 pass@k 指标（或其连续形式 maxg@k），定义为从策略中抽取 $k$ 个独立样本的最大期望奖励。

核心理论贡献

主要贡献在于推导出了针对 pass@k 目标及其梯度的新型、低方差、无偏估计量。这些估计量通过对来自 $n$ 个样本的 $\binom{n}{k}$ 个大小为 $k$ 的子集进行平均来实现。

1. 二值奖励 (第 2 节)：

   • 对于二值正确性标签（$f \in \{0, 1\}$），作者推导出了一个无偏估计量 $\rho(n, c, k)$ 用于表示 pass@k，其中 $c$ 是正确样本的数量。
   • 他们证明了仅当 $n \ge k$ 时才存在无偏估计量。
   • 随后推导出了相应的无偏梯度估计量，根据样本的正确性以及其被包含在成功子集中的组合概率来分配权重。

2. 连续奖励 (第 3 节)：

   • 该框架通过 maxg@k 目标 $E[\max_{i=1}^k g(x_i)]$ 将其推广到连续奖励 $g(x)$。
   • 通过假设奖励已排序（$g_1 \le g_2 \le \dots \le g_n$），作者利用二项式系数推导出了一个无偏估计量 $\rho^{(g)}(n, c, k)$。
   • 构建了一个无偏梯度估计量，其中样本的权重 $s_i$ 是奖励的线性组合，计算复杂度为高效的 $O(k + n \log n)$。

3. 方差缩减 (第 4 节)：

   • 为了在不引入偏差的情况下降低梯度估计的方差，作者提出了 留一法 (Leave-One-Out, LOO) 基准。
   • 他们引入了一个特定的基准 $s^{(loo-1)}_i$，该基准从大小为 $k$ 的子集的 maxg@k 中减去大小为 $k-1$ 的子集的期望 maxg@k（排除当前样本）。这推广了之前的工作（特别是 [TZSM25]），后者仅限于 $n=k$ 的情况。
   • 这种公式化方法使得即使在 $n > k$ 的情况下也能进行稳健优化，而 $n > k$ 是 LLM 训练中常见的采样场景。

实现

该方法被设计为任何策略梯度算法（如 PPO）中标准奖励变换的“即插即用”替代方案。给定一个包含 $n$ 个样本及其原始奖励的批次，在将奖励输入策略更新步骤之前，使用推导出的估计量（特别是提供的代码中的 sloo_minus_one 函数）对奖励进行变换。

关键结果

作者通过使用 GEMMA2 (2B 和 9B 参数) 和 LLAMA3.1 (8B 参数) 进行玩具实验和真实世界基准测试，验证了 PKPO 的有效性。

1. 玩具实验：

   • 在 1D 高斯策略实验中，PKPO 成功地使最优策略均值向更“风险偏好”的方向移动（针对更高的 $k$），允许模型探索具有较低概率但更高潜在回报的区域。
   • 方差分析证实，所提出的 sloo_minus_one 估计量与朴素划分或没有基准的估计量相比，显著降低了梯度方差。

2. 真实世界基准 (MATH, Coding, ARC-AGI-1)：

   • 选择性优化： 将优化目标 $k_{opt}$ 设置为等于评估指标 $k_{eval}$ 时，效果最佳。例如，优化 pass@8 与优化 pass@1 相比，能显著提高 pass@8 的评估得分。
   • 解决难题： 在极具挑战性的 ARC-AGI-1 基准测试中，传统的 pass@1 优化会陷入停滞，无法提高累积解决率。相比之下，使用 $k_{opt} > 1$ 的 PKPO “解锁”了学习过程，在所有 $k$ 下都实现了更高的解决率和 pass@k 得分。
   • 退火策略： 作者展示了在训练期间对 $k$ 进行退火（从高 $k_{opt}$ 开始以鼓励探索，然后降低到 $k_{opt}=1$ 以进行巩固）可以获得强大的 pass@1 性能，同时获得显著的 pass@k 增益，避免了典型的权衡问题（即提高 pass@k 会降低 pass@1）。
   • 对比： PKPO 优于 $n=k$ 的特例（如 [TZSM25]）以及熵正则化基准，特别是在需要关键探索的困难任务上。

重要性与主张

论文声称 PKPO 提供了一种稳健、高效且具有理论依据的方法，用于优化一组样本的集体效用，而非仅仅是平均值或单个最佳样本。

• 解锁探索： 作者认为，通过优先考虑联合效用（即 $k$ 个样本中至少有一个成功的概率）而非单个样本的效用，该方法改善了探索。这对于解决那些“信号”（正确的解）在搜索空间中非常稀疏的难题至关重要。
• 灵活性： 不同于以往仅限于 $k=n$（优化整个批次中的最佳者）或 $k=1$ 的方法，PKPO 能够针对任何任意的 $k \le n$ 进行稳健优化。
• 实际影响： 该方法允许从业者通过参数 $k_{opt}$ 来调节探索与利用之间的权衡，从而使模型能够解决那些通过标准 pass@1 优化无法触及的复杂推理和编程任务。

这项工作将自己定位为现有阶统计量策略梯度方法的泛化和方差缩减版本，为在需要深度探索的任务上扩展 LLM 的 RL 训练提供了一条实用的路径。