💬 NLP

On Advantage Estimates for Max@K Policy Gradients

本文介绍了 MaxPO，一种用于在具有可验证奖励的强化学习中优化 max@K 目标的全新策略梯度方法，该方法利用一种新颖的留一减二（Leave-Two-Out）基线来确保优势函数的中心化、降低梯度方差，并统一现有估计量，从而实现更有效的 LLM 后训练。

原作者： Shota Takashiro, Soichiro Nishimori, Paavo Parmas, Yongmin Kim, Kohsei Matsutani, Gouki Minegishi, Yusuke Iwasawa, Takeshi Kojima, Yutaka Matsuo

发布于 2026-06-05

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Shota Takashiro, Soichiro Nishimori, Paavo Parmas, Yongmin Kim, Kohsei Matsutani, Gouki Minegishi, Yusuke Iwasawa, Takeshi Kojima, Yutaka Matsuo

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你是一位正在试图训练一名学生解决极其困难数学题的老师。这个学生是一个人工智能（AI），而这个问题是一个“推理任务”。

在过去训练这些 AI 的方式中（强化学习），老师会要求学生尝试解决问题一次。如果答案错了，学生得不到任何反馈（奖励为零）；如果对了，他们会得到一颗小星星。问题在于，学生必须在巨大的可能性迷宫中猜中正确的路径，而获得小星星的机会非常渺茫，这导致学生经常陷入困境，不知道哪些猜测是“接近正确”的，哪些是“完全错误”的。

为了解决这个问题，研究人员开始了一种新的策略：“多次尝试”法。

老师不再要求学生只解决一次问题，而是要求他们在同一时间生成 K 个不同的解法。目标不仅仅是得到一个正确答案，而是要确保这 K 次尝试中至少有一个是正确的。这被称为优化 Max@K（或 Pass@K）。

旧有“多次尝试”法的问题

论文指出，虽然这种“多次尝试”的方法很棒，但用于教导 AI 如何改进的数学方法存在轻微的缺陷。

想象一下，你有一组 8 名学生（一个“批次”）正在尝试解决一个谜题。你想告诉每位学生他们表现得如何。

旧方法 (EI-only)： 它观察一个学生的答案，并将其与另外 7 名学生中的最佳答案进行比较。如果这个学生表现得比其他人好，他就会得到一个巨大的“做得好！”信号。如果表现得比其他人差，他就会得到一个“0”信号。
缺陷： 因为这个“做得好！”的信号始终是正向的（或为零），从未是负向的，所以老师实际上是在说：“你总是做得比平均水平更好！”这具有误导性。这就像一位教练，他只会说“太棒了！”而从不指出某个选手实际上低于团队的真实潜力。这创造了大量的“噪声”（方差），使得 AI 的学习过程变得不稳定且缓慢。

解决方案：“留一去二”基准线 (Leave-Two-Out Baseline)

论文作者提出了一种更聪明的评分方式，称之为 MaxPO (Max@K 策略优化)。

他们引入了一种新的评分规则，称为 “留一去二” (Leave-Two-Out, L2O) 基准线。它是这样运作的，让我们用一个简单的类比来说明：

想象你正在评判一场拥有 8 名选手的才艺表演。

旧方法： 为了评判选手 A，你将他与另外 7 人中的佼佼者进行比较。如果 A 是最优秀的，他会得到高分。如果不是，他会得到零分。这是有偏见的，因为“另外 7 人中的最佳”是一个随着房间里的人员变化而变化的移动目标。
新方法 (L2O)： 为了评判选手 A，你暂时将选手 A 和选手 B 从房间里移走。然后，你观察剩下的 6 个人，看看什么样的表现才是“公平”的平均水平。
- 你计算选手 A 相对于这个“公平”群体表现如何。
- 至关重要的是，通过移除两个人，你确保了“公平”的群体不会意外地包含你正在评判的对象（选手 A）或可能扭曲比较结果的特定“对手”（选手 B）。

为什么这更好？
这种方法确保了整个群体的“平均”得分正好为零。一些学生会得到正分（他们表现得比公平的平均水平好），而另一些人会得到负分（他们表现得比平均水平差）。

结果： AI 得到了一个更加清晰、噪声更小的信号。它确切地知道自己相对于一个公平的基准线处于什么位置，而不是仅仅被告知“你很棒”或“你不行”。

这篇论文的发现

研究人员通过两种方式测试了这种新的“留一去二”方法：

在简单游戏（多臂老虎机和迷宫）中： 他们展示了这种新方法在极大程度上降低了学习信号中的“噪声”（在某些情况下减少了高达 77% 的噪声）。这意味着 AI 的学习更加稳定，不会被随机波动所迷惑。
在真实的 AI 模型（大语言模型）中： 他们在大型语言模型（如 Llama 和 Qwen）尝试解决数学问题时测试了该方法。
- 结果： 使用这种新方法训练的 AI (MaxPO) 在允许尝试多次（例如 Pass@256）时，解决问题的成功率显著提高。
- 在 Qwen 模型上，成功率提升了 5.2%。
- 在 Llama 模型上，成功率提升了 2.4%。

大局观

把旧方法想象成一位过度乐观的教练，即使学生在挣扎，他也会给每个人发“做得好”贴纸。而新方法 (MaxPO) 则是一位使用严格、公平且平衡的评分系统的教练。通过消除“噪声”并确保得分以零为中心，AI 可以学习得更快、更有效，尤其是当目标是寻找多次尝试中的至少一个正确答案时。

论文得出结论，这种“留一去二”的方法是训练这类“多次尝试”任务的在数学上正确的做法，为未来的改进提供了一个统一且稳定的基础。

技术摘要：关于 Max@K 策略梯度的优势估计

问题陈述

带有可验证奖励（RLVR）的强化学习是后训练大语言模型（LLM）处理推理任务的主流范式。然而，这些设置通常依赖于稀疏的、基于结果的奖励（例如，二元的正确/错误信号），这使得在大型组合空间中的探索变得困难。虽然熵正则化和基于计数的奖励加成等方法存在，但它们引入了需要仔细调优的辅助目标。

一种互补的方法是直接优化推理时目标，例如 pass@K（ $K$ 个采样完成中至少有一个正确的概率）及其连续泛化形式 max@K（ $K$ 个样本中期望最大奖励）。最近的几项工作提出了针对这些目标的无偏策略梯度（PG）估计器。然而，这些估计器之间的关系仍不明确。具体而言，现有的估计器通常使用不同的信号、基线和归一化方式。本研究识别出的一个关键问题是，领先的估计器（例如基于期望改进（Expected Improvement, EI）的估计器）作为 PG 估计器是无偏的，但会产生非中心化的优势（non-centered advantages）。非中心化优势（均值不为零的优势）对应于被低估的基线，这会将优化导向次优动作并增加梯度方差。

方法论

本文提出了 MaxPO（Max@K 策略优化），该方法引入了 留二出（Leave-Two-Out, L2O） 基线，以修正 max@K 目标的优势估计。

1. 对现有估计器的分析

作者重新审视了先前工作（如 Walder 和 Karkhanis [54]）中使用的期望改进（EI）估计器。他们证明，虽然基于 EI 的估计器是无偏的 PG 估计器，但产生的优势信号 $s(a)$ 是严格非负的（由于 EI 定义中的 ReLU 操作）。因此， $E[s(a)] > 0$ ，这意味着优势不是中心化的。这种缺乏中心化的现象为方差缩减留下了空间。

2. 留二出 (L2O) 基线

为了在保持 PG 无偏性的同时实现精确的中心化优势，作者引入了 L2O 基线。

挑战： 标准的留一出（Leave-One-Out, L1O）基线在 max@K 设置中会失效，因为样本 $i$ 的奖励会隐式地影响其他样本 $j$ 的优势计算（因为 $i$ 在 max 操作中作为一个潜在的比较对象）。这种依赖性违反了无偏 PG 估计所需的独立性条件。
解决方案： L2O 基线 $b_{-i}^{L2O}$ 通过对所有其他样本 $j$ 的期望改进得分进行平均来计算，其中每个得分都是在排除目标样本 $i$ 和样本 $j$ 本身的比较集合上计算的。
特性：
- 独立性： 通过在所有 $j$ 的比较集中排除 $i$ ，基线 $b_{-i}^{L2O}$ 与奖励 $r_i$ 相互独立，确保了 PG 估计器保持无偏。
- 中心化： 作者证明了实现的批次平均优势对于每个采样的批次而言精确为零： $\frac{1}{B}\sum_{i=1}^B (s_i - b_{-i}^{L2O}) = 0$ 。
- 效率： 虽然朴素计算是 $O(B^3)$ ，但作者推导出了一个使用基于 ReLU 的公式实现的向量化版本，该版本能以 $O(B^2)$ 的时间复杂度计算 L2O 基线，使其能够应用于 GPU 加速的 LLM 训练。

3. 规范有限批次优势

论文推导出了 max@K 的规范有限批次优势，定义为 $u_i - v_i$ ：

$u_i$ ：包含响应 $i$ 的组的条件期望 max@K 值。
$v_i$ ：从排除 $i$ 的组中计算出的留一出期望 max@K 值。
作者表明，MaxPO 估计器（EI + L2O）精确地恢复了这一规范信号（ $s_i - b_{-i}^{L2O} = u_i - v_i$ ）。这提供了一个统一的视角来观察现有的估计器，将其分类为“非中心化”（如仅使用 EI/PKPO）和“（归一化）规范型”（如 MaxPO 以及某些解析形式的 pass@K 变体）。

4. 集成到 LLM 训练中

MaxPO 被集成到基于组的 RL 算法（如 GRPO）中。对于每个提示词（prompt），采样一组 $G$ 个输出。L2O 基线在每个组内计算，以生成序列级优势，随后这些优势被广播到 token 级的损失函数中。

核心贡献

识别非中心化优势： 论文指出，领先的 max@K PG 估计器（特别是基于 EI 的估计器）虽然作为梯度估计器是无偏的，但无法产生中心化的优势，从而导致次优的方差特性。
留二出 (L2O) 基线： 一种新型基线设计，在非线性 max@K 设置下既保留了 PG 无偏性，又确保了实现的批次优势为零均值。
规范统一： 推导了规范有限批次优势 ( $u_i - v_i$ ) 并证明了 MaxPO 恢复了该信号，为比较和分类现有的 pass@K/max@K 估计器提供了统一框架。
高效实现： 一个用于计算 L2O 基线的 $O(B^2)$ 向量化算法，使其能够实际应用于 LLM 后训练。

实验结果

作者通过受控的 Bandit/迷宫实验和 LLM 推理任务验证了该方法。

方差缩减： 在 Bandit 环境中，与原始 EI 估计器相比，L2O 基线显著降低了梯度方差，且这种效果随着动作空间大小的增加而增强。在 LLM 训练（Llama-3.2-3B-Instruct）中，该方法相比于 PKPO（即仅使用 EI 的基线）将梯度方差代理降低了 77.4%。
性能提升： 在五个数学推理基准测试（AIME24, AIME25, AMC23, MATH500, Minerva）上，MaxPO 优于包括 PKPO、GRPO 和 Entropy-Adv 在内的强基线。
- 在 Qwen2.5-Math-7B 上，MaxPO 比 PKPO 在任务平均 pass@256 上提升了 5.2%。
- 在 Llama-3.2-3B-Instruct 上，它比 PKPO 在 pass@256 上提升了 2.4%。
推理计算缩放： MaxPO 的性能增益在更高推理计算水平（较大的 $k$ 在 pass@k 中）时最为显著，这表明该方法有助于策略更好地利用额外的测试时样本。

重要性与声明

论文声称 MaxPO 为设计针对推理时目标（如 max@K）的策略梯度估计器提供了原则性的基础。通过识别规范优势形式并修正现有估计器中的偏差，这项工作提供了一种提高后训练推理模型效率和稳定性的方法。

作者谦虚地指出，其分析侧重于 i.i.d. 采样下的有限批次、同策略（on-policy）估计。他们承认关于离策（off-policy）训练、自适应采样以及相关生成设置的研究方向仍有待进一步探索。此外，他们强调，虽然该方法改进了优化过程，但更好地优化重复采样推理目标具有更广泛的影响，可能强化那些可能被误用于欺骗性协助或自动化问题解决的模型。