🤖 machine learning

PS-PPO: Prefix-Sampling PPO for Critic-Free RLHF

该论文提出了 PS-PPO，一种计算高效且无需评论家（critic-free）的 RLHF 方法，它通过仅对随机采样的轨迹前缀进行反向传播，并利用重要性采样权重来维持无偏梯度估计器，从而降低了训练成本和内存占用。

原作者： Doo Hwan Hwang, Kee-Eung Kim

发布于 2026-06-30

📖 1 分钟阅读☕ 轻松阅读

原作者： Doo Hwan Hwang, Kee-Eung Kim

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

想象一下，你正在教一名非常聪明的学生（一个人工智能）如何解决复杂的数学问题。这个学生在纸上写下他完整的思维过程，一步步地进行。在最后，你会检查最终答案。如果对了，你会给他一颗金星；如果错了，你会给他一个红色的 X。

问题：“全篇惩罚”
在目前的方法中（比如论文中提到的“无评论家”基准方法），当学生得到金星或红 X 时，老师会将整张纸视为同等重要。老师会重新阅读从第一句话到最后一句话的每一个字，试图弄清楚下次该如何改进。

论文指出，这是一种浪费时间。通常，当学生写到解题过程的中部时，结果是正确还是错误往往已经显而易见了。最后几句话通常只是“填充物”或重复内容。然而，即使如此，老师仍然强迫计算机每次都要为最后这几句话重新进行数学计算。这就像是为了决定是否喜欢一个故事，而去重新阅读一本你已经知道结局的书，这既耗费精力（计算能力）又占用内存。

解决方案：PS-PPO（“智能剪切”）
作者提出了一种新方法，称为 PS-PPO（前缀采样近端策略优化）。你可以把它想象成一种“智能剪切”策略。

与其每次都重新阅读整篇论文，老师使用一条特殊的规则来决定：“我到底需要读多少内容才能学到东西？”

“不确定性”计量器： 老师在学生写作的过程中观察他的作品。如果学生开始时有一个扎实的计划，老师就知道结果可能已经注定了。老师会说：“好吧，我不需要读剩下的部分了。”
随机截断： 老师随机选择一个停止阅读的位置（一个“截断点”）。也许在完成 30% 的内容后停止，也许在 70% 后停止。
公平性技巧（神奇秘诀）： 这是聪明之处。如果老师只读了前 30%，他可能会错过后面发生的某些重要事情。为了解决这个问题，老师使用了一个数学上的“公平性调整”。
- 如果他们提前停止，他们会给前 30% 的词汇赋予更高的权重。
- 如果他们读得较晚，则降低这些词汇的权重。
- 结果： 尽管他们只读了部分内容，但他们在许多、许多个案例中所学到的“平均课程”，与每次都阅读全文所学到的内容完全相同。

为什么这很重要
论文在难题（如高中竞赛题目）上测试了该方法。以下是他们的发现：

速度： 因为计算机停止了对句子结尾的计算，所以训练速度更快。他们观察到训练时间下降了约 33% 到 45%。
内存： 它使用了更少的计算机内存（RAM），这就像是做功时需要一张更小的书桌。
性能： 尽管阅读的内容变少了，但 AI 学习的效果与那些阅读全部内容的模型一样好。它获得了同样数量的金星。

“预算”类比
想象你每天有 100 个“思考标记（tokens）”的预算用于纠正你的学生。

旧方法： 你在每一张试卷上都花费所有的 100 个标记，即使前 20 个标记就已经足以看出错误。你会很快耗尽能量。
PS-PPO 方法： 你明智地使用这 100 个标记。对于简单的试卷，你只花 20 个标记；对于错误隐藏得很深的难题，你花 80 个标记。但因为有了“公平性调整”，你仍然能学到正确的教训。你在相同的时间内能处理更多的试卷，且不会损失质量。

总结
PS-PPO 是一种让 AI 模型更高效的方法。它意识到，对于冗长的复杂推理任务，故事的结尾往往不会增加新的信息。通过“剪切”掉故事的结尾，并利用数学调整来补偿教学过程，AI 能够以显著减少的计算能力和时间，达到同样优秀的学习效果。

技术摘要：PS-PPO：用于无评论员（Critic-Free）RLHF 的前缀采样 PPO

1. 问题陈述

基于人类反馈的强化学习（RLHF）是使大语言模型（LLM）与人类偏好对齐的标准方法。然而，使用 RLHF 训练 LLM 面临着显著的计算瓶ال瓶颈，特别是在无评论员（critic-free）方法（例如 GRPO、RLOO）中，这些方法避免了训练单独的价值网络。

在这些无评论员的方法中，单个标量奖励在生成完成后被分配给整个轨迹（completion）。随后，该奖励被均匀地广播到轨迹中的所有 token，以计算优势（advantages）。因此，策略更新需要针对每个 rollout 进行全轨迹反向传播。作者认为这是低效的，因为在许多推理任务（如逐步数学运算）中，中间的前缀往往已经包含了足以确定最终结果的充分信息。作者的实证分析（图 1）表明，轨迹的成功率往往在完成结束前就已趋于稳定，这意味着后缀部分的 token 携带了冗余的学习信号。因此，全轨迹更新会在这些冗余的后缀上浪费大量的计算资源和 GPU 显存。

2. 方法论：PS-PPO

作者提出了前缀采样近端策略优化（Prefix-Sampling PPO, PS-PPO），这是一种计算高效的无评论员方法，通过截断反向传播过程来利用时间上的冗余性，同时保持无偏的梯度估计量。

核心机制

PS-PPO 并非通过整个完成序列 $o = [o_1, \dots, o_T]$ 进行反向传播，而是为每个轨迹采样一个截断时间步 $H$ 。策略更新仅在当前前缀 $o_{\le H} = [o_1, \dots, o_H]$ 上进行计算。

为了确保这种截断不会引入偏差，PS-PPO 采用了重加权方案：

截断分布： 定义了一个以 prompt 为条件的分布 $\xi_{1:T}$ ，其中 $\xi_t = \Pr(H \ge t | x)$ 表示第 $t$ 个 token 被保留的概率。该分布是非增的（ $\xi_1 \ge \xi_2 \ge \dots \ge \xi_T$ ）。
无偏估计量： 截断轨迹的梯度通过逆包含概率 $1/\xi_t$ 进行重加权。所得估计量 $\hat{G}(\theta)$ 为：
$\hat{G}(\theta) := \frac{1}{K} \sum_{k=1}^K \sum_{t=1}^{H^{(k)}} \frac{1}{\xi_t} g_t^{(k)}(\theta)$
其中 $g_t^{(k)}$ 是逐时间步的梯度。该估计量在 $H$ 的采样上的期望值恢复了全轨迹更新。

优化截断分布

论文将 $\xi_{1:T}$ 的选择表述为一个凸优化问题。目标是在满足计算预算 $B$ （期望的反向传播 token 数）的条件下，最小化截断梯度估计量的方差。

方差代理： 作者推导出了重要性采样引起的方差的一个易于处理的代理，将其近似为 $\sum w_t(x) (1/\xi_t - 1)$ ，其中 $w_t(x)$ 代表第 $t$ 个时间步的更新重要性。
重要性代理： 由于计算精确的得分范数 $\|\nabla_\theta \log \pi_\theta\|^2$ $∥ \nabla_{θ} lo g π_{θ} ∥^{2}$ 需要进行完整反向传播（这会违背初衷），作者使用了一个基于输出头激活和奖励不确定性的仅前向代理。
- 奖励不确定性 ( $u_t(x)$ )： 通过给定前缀状态下的奖励方差来估计，该方差通过成功与失败的 rollout 之间下一 token 分布的距离来近似。
- 梯度范数代理 ( $\bar{\gamma}_t$ )： 从输出头梯度范数中导出。
优化： 该问题在满足单调性和预算约束的情况下，最小化 $\sum \frac{\bar{\gamma}_t(x,t) u_t(x)}{\xi_t}$ 。解法使用**相邻违反者消除（Pool Adjacent Violators, PAV）**算法来强制执行单调性约束。

算法流程

为每个 prompt 生成 $K$ 个 completion。
计算终止奖励并广播优势。
从 rollouts 中估计逐时间步的不确定性和梯度代理。
求解凸优化问题以获得最优的单调截断概率 $\xi^*_{1:T}$ 。
根据 $\xi^*$ 定义的分布为每个轨迹采样一个截断点 $H^{(k)}$ 。
6.仅对 $t \le H^{(k)}$ 的 token 进行反向传播，并使用 $1/\xi^*_t$ 对损失进行重加权。

3. 主要贡献

PS-PPO 框架： 一种新型的无评论员 RLHF 方法，它引入了以 prompt 为条件的随机截断反向传播过程，显著降低了计算和内存成本。
无偏截断估计器： 推导出了一个通过包含概率重加权的估计器，该估计器在期望上恢复了全序列广播更新，且不需要额外的 rollout 或辅助价值模型。
优化的截断策略： 一种确定“在哪里”进行截断的原则性方法，通过求解一个平衡方差减少与计算预算的凸优化问题，并利用前向代理来获取梯度重要性和奖励不确定性。
实证验证： 证明了 PS-PPO 在数学推理基准测试（MATH500, AMC 2023, Minerva Math, AIME 2024/2025）上达到了与强力无评论员基线（GRPO, DAPO, RLOO）相当的准确率，同时大幅减少了训练时间和峰值 GPU 显存。

4. 实验结果

作者使用 Llama-3.1-8B-Instruct 和 Qwen2.5-Math-7B 在数学推理基准测试上评估了 PS-PPO。

性能： PS-PPO（优化版）在所有基准测试中实现的 Pass@1 准确率与 DAPO 和 S-GRPO 等强基线相当或略好。例如，在 MATH500 上使用 Llama-3.1-8B 时，PS-PPO 达到了 47.6%，而 DAPO 为 46.8%。
效率：
- 训练时间： 与基线相比，PS-PPO 将每步训练时间减少了 33%–45%。这主要是由于减少了前向/反向成本，抵消了计算截断分布带来的开销。
- 显存： 由于减少了截断序列的激活存储，峰值 GPU 显存使用量降低了 15%–17%。
- 扩展性： 随着最大完成长度（ $T_{max}$ ）的增加，效率提升更加显著。在 $T_{max}=4096$ 时，PS-PPO 在更新阶段比 S-GRPO 快 2.8 倍，比 DAPO 快 3.3 倍。
消融实验：
- 预算 ( $B$ )： 预算 $B=128$ 个 token 提供了准确率与速度之间的最佳平衡。较小的预算会增加方差；较大的预算虽然能提高准确率，但计算成本收益递减。
- 截断策略： “优化”策略（求解凸问题）明显优于均匀、时间先验和启发式截断策略，证实了自适应选择信息丰富的序列前缀至关重要。
- Rollouts ( $K$ )： 增加 $K$ 可以提高稳定性和性能，但会增加每步的计算量。研究确定 $K=8$ 是一个实际的折中方案。

5. 重要性与声明

论文声称 PS-PPO 提供了一条在显著降低资源需求的情况下对齐大语言模型的可扩展路径。其核心意义在于观察到推理轨迹中的中间前缀通常包含足以确定最终结果的信息，使得全轨迹更新在计算上是冗余的。

通过通过随机前缀采样和重要性加权将学习信号与全序列长度解耦，PS-PPO 实现了：

降低硬件门槛： 更低的峰值显存和计算成本使得 RLHF 对学术界和资源受限的群体更加友好。
高效的长上下文训练： 该方法对于长推理轨迹（其全反向传播成本极高的情况）特别有效。

作者保持了谦逊的语气，指出 PS-PPO 是一种优化方法，旨在提高效率，但它本身并不决定模型的安全性或行为；这些仍取决于奖励规范、训练数据和部署保障措施。该方法在期望上保持了全序列目标的理论特性，确保了效率的提升不会以产生有偏学习为代价。