PS-PPO 阅读笔记:跳过轨迹里「已经确定」的部分,但不破坏 PPO 的数学

关于本文范围的说明: 本笔记从零开始搭建所有前置知识——RLHF 训练循环、PPO 的裁剪目标函数、无 critic 的优势估计(GRPO 一类的做法)、重要性采样与”纳入概率”重加权、以及拉格朗日松弛下的凸对偶——然后完整走一遍 PS-PPO 的推导链条:无偏性证明、方差代理量的推导、只用前向传播的 score-norm 代理、奖励不确定性的上界、以及用 Pool-Adjacent-Violators(PAV)算法求解的单调预算方案,每一步都配有手算数值例子。阅读本文只需要基本的概率论知识和对梯度下降的熟悉程度。

笔记日期: 2026-07-21 笔记作者: Zhongzhu Zhou 论文标题: PS-PPO: Prefix-Sampling PPO for Critic-Free RLHF 作者: Doo Hwan Hwang, Kee-Eung Kim arXiv: 2606.29758 状态: KAIST(Kim Jaechul 人工智能研究生院);实现代码发布于 github.com/doohwan383/PS-PPO,基于 Hugging Face 的 Open-R1 代码库

一句话总结

现在主流的、不用 critic 的 RLHF 方法——GRPO、RLOO、DAPO 以及它们的各种变体——都共享一个隐藏的低效之处:它们对整段生成的回复只计算一个标量奖励,然后把这个同一个标量当作回复中每一个 token 的”优势”广播下去。这就导致每次梯度更新都必须对整段序列做完整的前向和反向传播,不管回复到底有多长。PS-PPO 从一个简单的经验观察出发——在逐步数学推理任务中,一段回复最终是否正确,往往在只看到前 25%-40% 的 token 时就已经高度可预测了(图 1)——并进一步追问:能不能让更新本身也提前结束,而不只是奖励计算提前结束?这个想法最朴素的版本(固定截断每段回复,或者随机截断)会引入偏差:你实际上在优化一个不同的、错误的目标。PS-PPO 真正的贡献在于给出一个有理论保证的、无偏的做法:把截断长度当作一个随机变量,赋予它一个精心设计的、依赖于 prompt 的”生存分布”,推导出一个重要性加权修正项,可以证明其恢复了与全序列更新完全相同的期望梯度;然后把”在固定算力预算下,选择哪个生存分布能让这次修正引入的额外方差最小”这个问题,构造成一个凸优化问题,并给出通过 Pool-Adjacent-Violators 算法求解的、干净的(分块)闭式解。实际效果是直接测量出来的,而不是空口断言的:相比 DAPO / S-GRPO 一类的无 critic 基线,梯度更新阶段的墙钟时间减少 33%-45%,GPU 峰值显存减少 15%-17%,同时在 MATH500、AMC23、CollegeMath、MinervaMath、AIME24、AIME25 六个基准、Llama-3.1-8B-Instruct 与 Qwen2.5-Math-7B 两个骨干模型上,pass@1 准确率在统计意义上没有显著差异(表 2);而且回复越长,这个效率优势就越明显(在最大长度 $T_{max}=4096$ 时,比只做”掩码”的基线快 2.8-3.3 倍,表 3)。

核心要点速览

• 根本问题: 无 critic 的 RLHF 把一个标量奖励 $R(x,o)-b(x)$ 广播到回复的每一个 token 上;由此产生的策略梯度更新,即便许多轨迹的结果其实早就已经确定,仍然需要对每一个 token 做完整的前向/反向传播。
• 实证动机(图 1): 在 Qwen2.5-Math-7B 上,以 prefix 为条件的成功率在 AIME 2024 上只用到完整回复长度的 25.0% 就跨过了 $\geq 0.95$ 的阈值,在 MATH-500 上是 38.2%——也就是说,回复的很大一部分”尾巴”其实对最终奖励几乎不再提供额外信息。
• 核心思路: 不是确定性地或者用启发式方法截断,而是从一个”生存分布” $\xi_1\geq\xi_2\geq\cdots\geq\xi_T$(其中 $\xi_t=\Pr(H\geq t)$)中为每段回复随机采样一个截断时间步 $H$,只对 $t\leq H$ 的 token 做反向传播,并把每个保留 token 的梯度按 $1/\xi_t$ 重新加权。这样可以精确恢复全序列的期望梯度(附录 B,下文完整复现推导过程)——没有偏差,纯粹是方差与算力之间的权衡。
• 这个设计问题本质上是一个凸优化问题,而不是一个启发式规则: 在满足算力预算约束 $\sum_t t(\xi_t-\xi_{t+1})=B$ 和单调性约束 $1\geq\xi_1\geq\cdots\geq\xi_T>0$(单调性是必须的,因为 $H$ 必须表现得像一个真实的”生存时间”)的前提下,最小化一个可处理的、基于协方差迹的方差代理量。不考虑单调性约束时的解有一个干净的闭式 $\xi_t\propto\sqrt{w_t}$;要满足单调性约束,则需要借用等式回归(isotonic regression)里的 Pool-Adjacent-Violators(PAV)算法。
• 一个让权重计算变得便宜的巧妙技巧: 方差代理量所需要的每步权重 $w_t^\theta(x)$,正比于平方 score-function 范数 $\|\nabla_\theta\log\pi_\theta(o_t|s_t)\|^2$,而这通常需要一次完整的反向传播才能算出来——恰好是 PS-PPO 想要避免的开销。论文证明这可以用一个只依赖前向传播、闭式可算的代理量来近似(附录 D),只需要输出层的隐藏状态和 softmax 概率,并且在实证上与真实的全参数 score norm 的皮尔逊相关系数达到 $r=0.9968$(图 4)。
• 第二个巧妙技巧,针对优势不确定性项: 由于奖励往往是二值的(正确/错误),论文只用同一批 rollout 里已经有的下一个 token 的分布统计量,推导出了 $\mathrm{Var}(R\mid s_t)$ 的一个可证明的上界(附录 E)——不需要额外的 rollout,也不需要辅助的价值网络。
• 效率(表 1、图 2): 在默认预算 $B=128$ 下,PS-PPO 每步训练总时间为 $1.77\pm0.02$ 秒,而 S-GRPO 是 $2.66\pm0.01$ 秒,DAPO 是 $3.23\pm0.01$ 秒,带 forking-token 的 DAPO 是 $3.25\pm0.01$ 秒——减少了 33%-45%,即便已经算上了计算截断分布本身多花的 $0.43\pm0.02$ 秒开销。GPU 峰值显存下降 15%-17%。
• 只做”掩码”的基线(S-GRPO、带 forking-token 的 DAPO)其实并没有真正省下算力,因为它们只是把某些 token 的损失清零,但前向/反向传播仍然要跑过整个序列——PS-PPO 是这个对比组里第一个真正缩短了计算图本身的方法(图 2(d))。
• 准确率没有被牺牲: 在两个骨干模型上,PS-PPO(Optimized)在 6 个基准中的 5 个上都能打平或略微超过 DAPO(表 2)——被优化过的截断分布不只是”更便宜”,在相同的期望算力预算下,它在分配算力这件事上确实做得比均匀或启发式的替代方案更好(论文的表 4,下文会讨论)。
• 效率优势会随着回复长度增长而复合放大: 在 $T_{max}=4096$ 个 token 时,PS-PPO 每次更新比 S-GRPO 快 2.8 倍,比 DAPO 快 3.3 倍(表 3)——随着长思维链训练成为常态,这一点会越来越重要。
• 论文明确指出的边界条件: 所有主要实验都使用二值的正确性奖励;连续奖励的 RLHF 实验(附录 G,HH-RLHF/IMDB)也做了,但只是用一个启发式的 sigmoid 二值化去喂给同一个不确定性估计器——这是在一层近似之上又叠加了一层近似。

前置知识:你需要先了解什么

从语言模型到 RLHF 策略

一个自回归语言模型 $\pi_\theta$,逐个 token 地定义了”在已生成内容的条件下,下一个 token”的概率分布。在 RLHF 的设定里,我们把这个模型当作一个”策略”:给定一个从某个 prompt 分布 $p_Q$ 采样出来的 prompt $x$,策略逐个 token 采样出一段回复 $o=[o_1,\dots,o_T]$,其中每个 token $o_t$ 都从 $\pi_\theta(\cdot\mid s_t)$ 中采样,$s_t=[x,o_1,\dots,o_{t-1}]$ 是”状态”——也就是 prompt 加上目前为止生成的所有内容。整段回复生成完之后,一个奖励函数 $R(x,o)$ 给整段回复打一个标量分数——对于数学推理来说,通常是最终答案框里的内容与标准答案匹配就是 $R=1$,不匹配就是 $R=0$。

训练的目标是调整 $\theta$,使得 $\mathbb{E}_{x\sim p_Q, o\sim\pi_\theta}[R(x,o)]$ 增大。这是一个标准的强化学习目标,但相比经典 RL 有一个不寻常的结构:“episode” 就是一次完整的文本生成,“奖励” 通常只在最后到达一次(稀疏、终局奖励),而每一步的”动作空间”是整个模型词表(常常有 10 万以上个 token)。

PPO 的裁剪代理目标函数

近端策略优化(PPO)是把标量奖励信号转化为稳定参数更新的标准策略梯度算法。它的核心对象是逐 token 的重要性比率:

$$\rho_t(\theta) = \frac{\pi_\theta(o_t\mid s_t)}{\pi_{\theta_{\text{old}}}(o_t\mid s_t)},$$

它衡量的是:相比于实际生成这段轨迹的旧策略 $\pi_{\theta_{\text{old}}}$,当前策略 $\pi_\theta$ 生成同一个 token $o_t$ 的可能性变大了多少(或变小了多少)。PPO 的训练目标对这个比率做裁剪,防止单次更新把策略推得离生成数据的分布太远:

$$J_{\text{PPO}}(\theta) = \mathbb{E}_{x\sim p_Q,\, o\sim\pi_{\theta_{\text{old}}}(\cdot\mid x)}\left[\sum_{t=1}^{T}\min\Big(\rho_t(\theta)\hat A_t,\ \mathrm{clip}(\rho_t(\theta), 1-\epsilon, 1+\epsilon)\hat A_t\Big)\right],$$

其中 $\hat A_t$ 是第 $t$ 步的优势估计(衡量这个动作比”平均水平”好多少),$\epsilon$ 是裁剪范围(通常取 0.1-0.2)。裁剪的直觉是:如果新策略把某个 token 的概率抬得比旧策略高很多($\rho_t\gg 1$),且这个 token 的优势为正,那么目标函数会把这次更新的收益”封顶”,防止策略在一步内跳得太远、偏离生成数据的分布,导致后续估计失真。

无 critic 的优势估计:以 GRPO 为例

PPO 原本需要一个学习出来的价值函数(critic)来估计 $\hat A_t$,但训练一个 LLM 规模的 critic 本身就是一笔不小的算力和显存开销,而且经常是训练不稳定的主要来源。GRPO(Group Relative Policy Optimization)一类的方法完全绕开了这一点:对同一个 prompt 采样 $K$ 段完整回复,计算它们各自的终局奖励 $R^{(1)},\dots,R^{(K)}$,然后用组内均值 $b(x)=\frac{1}{K}\sum_k R^{(k)}$ 作为基线,把每段回复的优势定义为 $\hat A^{(k)} = R^{(k)}-b(x)$(有时再除以组内标准差做归一化)。关键之处在于:这个优势值对同一段回复里的每一个 token 都是同一个常数——这就是本文反复提到的”广播式优势”。这个设计选择本身是合理的(不需要额外训练一个 critic,方差通常也可控),但它的代价正是 PS-PPO 想要解决的问题:既然 $\hat A_t$ 对所有 $t$ 都相同,那”梯度到底该在哪个 token 上花更多算力”这件事,原始的 GRPO 完全没有建模。

重要性采样与”纳入概率”重加权:直觉铺垫

在深入 PS-PPO 的具体推导之前,有必要先建立一个更一般的直觉:如果你只随机保留一个总体中的一部分个体(而不是保留全部),但你想让保留下来的这部分个体的加权平均,恰好等于对全部个体求平均的结果,你需要怎么做?答案是经典的”逆概率加权”(inverse-probability weighting):如果个体 $i$ 被保留的概率是 $\pi_i$,那么把它保留下来的贡献除以 $\pi_i$ 再求和,期望上就精确等于对全体求和的结果——因为 $\mathbb{E}[\mathbb{1}\{i \text{ 被保留}\}/\pi_i] = \pi_i/\pi_i = 1$。这正是 PS-PPO 用 $1/\xi_t$ 重新加权保留 token 的数学基础,只不过这里”是否保留”的判定不是对每个 token 独立做的,而是由一个统一的截断时间点 $H$ 决定的——所有 $t\leq H$ 的 token 一起被保留、一起被重新加权。

PS-PPO 瞄准的计算瓶颈,用具体数字说清楚

考虑一个具体例子:一段 $T=2048$ token 的数学推理回复,在 8B 参数模型上做一次标准的 GRPO 更新。前向传播和反向传播的算力开销都大致随序列长度线性增长(注意力机制以及每层后续计算都要覆盖全部 $T$ 个位置),所以不管这段推理”实际上”在第 500 个 token 左右就已经能大概率确定对错,GRPO 仍然要为全部 2048 个 token 花费同样的前向/反向算力。图 1 的经验发现正是量化了这种浪费的规模:在 AIME 2024 上,只用了 25% 的长度就已经能以 $\geq 95\%$ 的置信度预测最终成功与否;在 MATH-500 上是 38%。如果能把反向传播的算力集中花在真正”决定成败”的那部分 token 上,理论上就能省下大部分被浪费掉的算力——这正是 PS-PPO 想要形式化并证明安全性的直觉。

flowchart LR
    P["prompt x, rollout 生成 tokens o_1 ... o_T"] --> D["从生存分布 xi_1:T x 中采样截断点 H"]
    D --> K1["t = 1..H 的 token: 保留, 梯度按 1/xi_t 重加权"]
    D --> K2["t = H+1..T 的 token: 丢弃, 不花前向/反向算力"]
    K1 --> B["截断且重加权的梯度求和 -> 输入 PPO 更新 (对 H 求期望后无偏)"]

图:PS-PPO 随机截断机制下,单段回复的数据流——截断点 $H$ 对每段回复、每次训练步都独立重新采样,所以同一个 batch 里不同回复会被截断在不同的长度,而对保留 token 做的 $1/\xi_t$ 重加权,正是让丢弃”尾巴”之后仍能恢复全序列期望的关键。

理论第一部分:推导无偏的截断估计量

把全序列更新写成一个求和

先写出无 critic 策略梯度更新的完整形式,对每个 prompt 的 $K$ 段回复求和:

$$g_t^{(k)}(\theta) := \hat A_t^{(k)}\,\nabla_\theta\log\pi_\theta\big(o_t^{(k)}\mid s_t^{(k)}\big), \qquad G(\theta) := \frac{1}{K}\sum_{k=1}^{K}\sum_{t=1}^{T}g_t^{(k)}(\theta).$$

这正是 GRPO/RLOO/DAPO 已经在计算的东西:对每一个 token,把它的逐 token score function(对数概率的梯度)乘以广播下来的优势值,对组内所有回复、所有 token 求和再平均。$G(\theta)$ 就是 PS-PPO 想要用更少算力去近似的”全序列真值”梯度方向。

引入随机截断点并重加权

现在假设每段回复 $k$ 都得到一个独立采样的随机截断点 $H^{(k)}$,只对 $t\leq H^{(k)}$ 的 token 做反向传播。定义生存概率 $\xi_t := \Pr(H\geq t\mid x)$——也就是”截断点还没到第 $t$ 步”的概率,即第 $t$ 个 token 被保留的概率。注意按约定 $\xi_1=1$(至少保留 1 个 token 的前缀),且 $\xi_1\geq\xi_2\geq\cdots\geq\xi_T$(生存函数必然是非增的——你不可能随着时间推移变得”更容易存活”)。朴素的截断估计量 $\frac{1}{K}\sum_k\sum_{t\leq H^{(k)}} g_t^{(k)}(\theta)$ 是有偏的——它系统性地低估了 $G(\theta)$,因为它悄悄丢掉了截断点之后那些(期望上可能非零的)贡献。修正方案是重加权估计量:

$$\widehat G(\theta) := \frac{1}{K}\sum_{k=1}^{K}\sum_{t=1}^{H^{(k)}}\frac{1}{\xi_t}\,g_t^{(k)}(\theta).$$

逐步证明无偏性

这是论文两个核心理论结果之一(附录 B),值得完整走一遍推导,而不是直接相信结论,因为整个方法的正确性都建立在这个证明上。定义纳入指示函数 $I_t^{(k)} := \mathbb{1}\{t\leq H^{(k)}\}$,这样上式可以改写成对所有 $t=1,\dots,T$ 求和(而不只是到截断点为止),因为 $t>H^{(k)}$ 的项自动有 $I_t^{(k)}=0$:

$$\widehat G(\theta) = \frac{1}{K}\sum_{k=1}^{K}\sum_{t=1}^{T}\frac{I_t^{(k)}}{\xi_t}\,g_t^{(k)}(\theta).$$

第一步——纳入指示函数的期望恰好是 $\xi_t$,这是定义使然: $\mathbb{E}[I_t^{(k)}\mid x] = \Pr(H^{(k)}\geq t\mid x) = \xi_t$。这只是重述了 $\xi_t$ 作为生存概率的定义,还不是证明,但它是后面一切的立足点。

第二步——先固定 rollout,只让截断点保持随机: 固定 prompt $x$ 和 $K$ 段已经采样出来的轨迹 $o^{(1)},\dots,o^{(K)}$(也就是对除了截断点 $H^{(k)}$ 之外的一切都取条件——截断点是在 rollout 完成之后独立采样出来的,其分布可以依赖 $x$,但不依赖某段具体轨迹实际采样到的 token,这种独立性是算法本身保证的,因为 $\xi_{1:T}(x)$ 是从 prompt 级和 batch 级的统计量算出来的,而不是从某段轨迹超出该点之后实际采样到的具体内容算出来的)。在这个条件下,$g_t^{(k)}(\theta)$ 是一个固定的、非随机的量,式子里唯一剩下的随机性来自 $I_t^{(k)}$:

$$\mathbb{E}\big[\widehat G(\theta)\mid x, o_{1:K}\big] = \frac{1}{K}\sum_{k=1}^{K}\sum_{t=1}^{T}\mathbb{E}\left[\frac{I_t^{(k)}}{\xi_t}\ \middle|\ x, o_{1:K}\right]g_t^{(k)}(\theta) = \frac{1}{K}\sum_{k=1}^{K}\sum_{t=1}^{T}\frac{\xi_t}{\xi_t}\,g_t^{(k)}(\theta) = \frac{1}{K}\sum_{k=1}^{K}\sum_{t=1}^{T}g_t^{(k)}(\theta) = G(\theta).$$

中间这一步是关键:$g_t^{(k)}(\theta)$ 在这个条件下是固定的,可以从条件期望里提出来,剩下 $\mathbb{E}[I_t^{(k)}/\xi_t\mid x,o_{1:K}] = \xi_t/\xi_t = 1$——每一项都精确地坍缩回 $g_t^{(k)}(\theta)$ 本身,求和之后精确恢复 $G(\theta)$,不是近似恢复。

第三步——再对 rollout 本身取期望: 既然上一步的结论对任意一次具体的轨迹实现都成立(这只是关于截断点随机性的一个陈述,在任意一组固定轨迹的条件下都成立),对两边再取一次关于 rollout 采样随机性的期望,就得到 $\mathbb{E}[\widehat G(\theta)\mid x] = \mathbb{E}[G(\theta)\mid x] = G(\theta)$(因为 $G(\theta)$ 本身根本不依赖截断点的随机性)。证明到此完成:对任意合法的生存概率选择 $\xi_{1:T}$,截断且重加权的估计量,其期望都与全序列估计量完全相同——$\xi_{1:T}$ 的选择只会影响方差,永远不会引入偏差。

这在实践中为什么重要?因为这意味着 PS-PPO 并不是在”更快”和”正确”之间做权衡——它是在”更快”和”方差”之间做权衡,而这是一个本质上更温和的权衡,因为方差可以通过精心选择 $\xi_{1:T}$ 来控制(这正是下一节要做的事),而偏差一旦引入,通常事后是没法修正的。

理论第二部分:把截断分布的设计构造成一个凸优化问题

为什么不能随便选一个 $\xi_t$?

上面的无偏性证明对任意单调、合法的生存序列 $\xi_{1:T}$ 都成立——包括像 $\xi_t=1$ 对所有 $t$(完全不截断,没有任何节省)这样糟糕的选择,或者 $\xi_t$ 衰减得极快、几乎什么都不保留(算力节省巨大,但 $1/\xi_t$ 重加权在极少数保留下来的 token 上会爆炸,方差极大)。论文真正的技术贡献是恰当地选择 $\xi_{1:T}$——具体来说,是在固定花费某个期望算力量的约束下,最小化截断引入的额外方差。

推导可处理的方差代理量(论文中的公式 4)

$\widehat G(\theta)$ 由截断随机性 $H$ 引入的精确方差(即 $\mathrm{Var}_H(\widehat G\mid x) \equiv \mathrm{tr}\,\mathrm{Cov}_H(\widehat G\mid x)$,协方差矩阵的一个迹/对角线标量摘要),原则上需要考虑跨时间步的相关性,因为同一段轨迹里相邻位置的 token 是自回归相关的。论文的做法是只保留对角线(逐时间步)方差项,丢弃跨时间步的协方差——这是一个刻意的、明确承认的近似,理由是保留交叉项会让设计变量 $\xi_t$ 跨时间步耦合起来,破坏掉让这个优化问题能有闭式解的干净结构。

完整推导如下(附录 C),每一步中间过程都展开:

第一步。 利用 $\mathbb{E}_H[\mathbb{1}\{H^{(k)}\geq t\}/\xi_t\mid x]=1$(上面已证明)这一事实,把 $\widehat G(\theta) - \mathbb{E}_H[\widehat G(\theta)\mid x]$ 显式写出来:

$$\widehat G(\theta) - \mathbb{E}_H[\widehat G(\theta)\mid x] = \frac{1}{K}\sum_{k=1}^{K}\sum_{t=1}^{T}\left(\frac{\mathbb{1}\{H^{(k)}\geq t\}}{\xi_t} - 1\right)g_t^{(k)}(\theta).$$

第二步。 因为不同回复 $k$ 之间的截断点是独立采样的,跨回复的协方差项会消失,基于协方差迹的代理量就变成:

$$\mathrm{Var}_H(\widehat G(\theta)\mid x) = \frac{1}{K}\sum_{t=1}^{T}\mathbb{E}_H\left[\left(\frac{\mathbb{1}\{H\geq t\}}{\xi_t}-1\right)^2\Big|x\right]\|g_t(\theta)\|^2,$$

这里用一个通用的 $(H,g_t)$ 代表任意一段回复(为了可读性省去了上标 $(k)$),而被丢弃的正是同一条轨迹内部跨时间步的协方差项。

第三步。 显式算出内层期望。由于 $\mathbb{1}\{H\geq t\}$ 是一个 Bernoulli$(\xi_t)$ 变量:

$$\mathbb{E}_H\left[\left(\frac{\mathbb{1}\{H\geq t\}}{\xi_t}-1\right)^2\Big|x\right] = \xi_t\left(\frac{1}{\xi_t}-1\right)^2 + (1-\xi_t)\cdot 1^2 = \frac{1}{\xi_t} - 1,$$

这是直接展开一个两种取值的期望得到的:以概率 $\xi_t$ 取值 $(1/\xi_t-1)^2$,以概率 $1-\xi_t$ 指示函数为 0、取值 $(0-1)^2=1$;代数化简 $\xi_t(1/\xi_t-1)^2+(1-\xi_t) = \xi_t\cdot\frac{(1-\xi_t)^2}{\xi_t^2}+(1-\xi_t) = \frac{(1-\xi_t)^2}{\xi_t}+(1-\xi_t) = (1-\xi_t)\left[\frac{1-\xi_t}{\xi_t}+1\right] = (1-\xi_t)\cdot\frac{1}{\xi_t} = \frac{1-\xi_t}{\xi_t} = \frac{1}{\xi_t}-1$ 确认了这个简化形式。

第四步。 代回原式,并对 rollout 取期望,定义 $w_t^\theta(x) := \mathbb{E}[\|g_t(\theta)\|^2\mid x] = \mathbb{E}[\hat A_t^2\|\nabla_\theta\log\pi_\theta(o_t\mid s_t)\|^2\mid x]$(第 $t$ 步梯度贡献的期望平方范数,涵盖采样到哪个 token、优势值最终是多少的随机性),就得到论文中的公式 4:

$$\mathbb{E}\big[\mathrm{tr}\,\mathrm{Cov}_H(\widehat G(\theta)\mid x)\big] \approx \frac{1}{K}\sum_{t=1}^{T} w_t^\theta(x)\left(\frac{1}{\xi_t}-1\right).$$

直觉解读: 截断引入的额外方差,是对每个时间步”这一步的梯度有多重要”($w_t^\theta(x)$)乘以”以较低概率被保留、需要承受的重加权惩罚”($1/\xi_t - 1$,当 $\xi_t\to 0$ 时爆炸,$\xi_t\to 1$ 时趋于 0)求和。这立刻指向了正确的设计原则:给重要梯度($w_t^\theta(x)$ 大)的时间步以更高的纳入概率 $\xi_t$,让不重要梯度的时间步可以有较低的纳入概率——这恰好与图 1 的经验曲线相吻合,因为一段已经变得可预测的轨迹,其后期 token 的 $w_t^\theta(x)$ 应该很小。

完整设计问题及其闭式解

在算力预算约束 $\sum_{t=1}^{T} t(\xi_t-\xi_{t+1}) = B$(这个表达式恰好就是 $\mathbb{E}[H]$,即期望保留的前缀长度——对非负整数随机变量有一个标准恒等式 $\mathbb{E}[H]=\sum_t \Pr(H\geq t)=\sum_t\xi_t$,经过 Abel 分部求和重排后就变成了这里用到的”望远镜求和”形式)和生存单调性约束 $1\geq\xi_1\geq\cdots\geq\xi_T>0$ 下最小化公式 4,等价于以下更简单的问题:

$$\min_{\xi_{1:T}} \sum_{t=1}^{T}\frac{w_t^\theta(x)}{\xi_t} \quad\text{s.t.}\quad \sum_{t=1}^{T}t(\xi_t-\xi_{t+1})=B,\quad \xi_{T+1}:=0,\quad 1\geq\xi_1\geq\cdots\geq\xi_T>0.$$

先求解放松掉单调性约束的版本。 暂时去掉单调性约束,考虑更简单的等式约束问题 $\min_{\xi_t>0}\sum_t w_t/\xi_t$,约束是 $\sum_t \xi_t = B$(用预算约束更简单的加法形式来近似——论文附录 F 直接用的就是这个等价形式)。写出拉格朗日函数:

$$\mathcal{L}(\xi,\lambda) = \sum_{t=1}^{T}\frac{w_t}{\xi_t} + \lambda\left(\sum_{t=1}^{T}\xi_t - B\right).$$

对单个 $\xi_t$ 求偏导并令其为零:

$$\frac{\partial\mathcal{L}}{\partial\xi_t} = -\frac{w_t}{\xi_t^2} + \lambda = 0 \implies \xi_t = \sqrt{\frac{w_t}{\lambda}}.$$

把这个形式代回预算约束 $\sum_t \xi_t = B$,得到 $\sum_t\sqrt{w_t/\lambda} = B$,即 $\frac{1}{\sqrt\lambda}=\frac{B}{\sum_j\sqrt{w_j}}$,于是不考虑单调性约束时的闭式解是:

$$\xi_t^\star = B\cdot\frac{\sqrt{w_t}}{\sum_{j=1}^{T}\sqrt{w_j}}.$$

直觉解读: 第 $t$ 步的最优纳入概率,正比于该时间步梯度重要性权重的平方根——不是线性正比。这个平方根关系直接来自目标函数中的 $1/\xi_t$ 项(对小 $\xi_t$ 的双曲惩罚)与线性预算约束之间的平衡——这正是分层抽样理论(Neyman 分配)里,凡是要在线性预算约束下最小化 $w_i/\xi_i$ 型求和目标时,都会出现的同一条平方根分配规则。

用分块合并恢复单调性。 上面的闭式解通常不会自动产生非增序列——$w_t$(梯度重要性)可以随时间步上下波动,所以 $\sqrt{w_t}$ 也会,但 $\xi_t$ 作为生存概率必须是非增的。论文用等式回归问题的标准解法来处理:把 $\{1,\dots,T\}$ 划分成若干连续的块,强制 $\xi_t$ 在每个块内保持常数,然后在块的层面求解同一个拉格朗日问题。对一个长度 $L_m=j_m-i_m+1$、聚合权重 $W_m=\sum_{t=i_m}^{j_m}w_t$ 的块 $[i_m,j_m]$,同样的推导(把”逐时间步”换成”逐块”)给出:

$$\xi_m^\star = B\cdot\frac{\sqrt{W_m/L_m}}{\sum_{j=1}^{M}\sqrt{L_jW_j}}, \qquad s_m := \sqrt{W_m/L_m}.$$

如果块分数 $s_1\geq s_2\geq\cdots\geq s_M$ 本身已经是非增的,那么这个分块解就已经是正确的单调全局最优解。如果不是——某个较晚的块分数反而更高,这就需要一个非单调的生存序列——Pool-Adjacent-Violators(PAV)算法就会把两个违反单调性的相邻块合并成一个更大的块(重新计算合并后块的 $W$ 和 $L$),重复这个过程直到块分数序列变成非增为止。PAV 是等式回归里一个广为人知的、精确最优、$O(T)$ 摊销复杂度的算法,所以这一步相对于前向/反向传播节省下来的算力来说,开销可以忽略不计。

完整”预算约束+单调截断”计算的伪代码:

输入: 逐时间步权重 w_1, ..., w_T (来自前向代理量, 见下文);
      算力预算 B
1. 计算无约束分数: score_t = sqrt(w_t), 对每个 t = 1..T
2. 把每个时间步初始化为独立的单点块: blocks = [{t}: t=1..T]
   块分数 block_score[t] = score_t, 块长度 block_length[t] = 1, 块权重 block_weight[t] = w_t
3. 当存在相邻块 (m, m+1) 满足 block_score[m] < block_score[m+1] 时:
       合并块 m 和 m+1 为单个新块 m':
           block_length[m'] = block_length[m] + block_length[m+1]
           block_weight[m'] = block_weight[m] + block_weight[m+1]
           block_score[m']  = sqrt(block_weight[m'] / block_length[m'])
       (这正是 Pool-Adjacent-Violators 的合并步骤)
4. 收敛后, block_score 在最终的块序列上是非增的
5. 对块 j 内的每个 t, 赋值 xi_t = B * block_score[block(t)] / sum_j(block_length[j] * block_score[j])
   这就是用来采样截断点的最终 xi_t^*
输出: 单调的截断概率 xi_1 >= xi_2 >= ... >= xi_T > 0

为什么用平方根分配,而不是线性或均等分配?一个数值对比

值得用实际数字验证一下,为什么推导出的结果是 $\xi_t\propto\sqrt{w_t}$,而不是听起来更直觉的 $\xi_t\propto w_t$(线性分配),或者 $\xi_t=B/T$ 对所有 $t$(均等分配)。取一个只有两个时间步的玩具例子,$w_1=9$(非常重要)、$w_2=1$(不重要),预算 $B=1$(即两个时间步加起来,期望上只能负担得起完整保留 1 个 token)。

均等分配 给出 $\xi_1=\xi_2=0.5$,不管重要性高低都平均花费预算——显然浪费,因为 $w_1$ 比 $w_2$ 重要 9 倍,却没有获得任何额外的纳入概率。

线性分配($\xi_t\propto w_t$)给出 $\xi_1 = 1\cdot\frac{9}{10}=0.9$,$\xi_2=1\cdot\frac{1}{10}=0.1$。代入真实要最小化的目标函数,$\sum_t w_t/\xi_t = 9/0.9 + 1/0.1 = 10+10=20$。

平方根分配(实际推导出的最优解)给出 $\xi_1=1\cdot\frac{3}{4}=0.75$,$\xi_2=1\cdot\frac{1}{4}=0.25$(因为 $\sqrt{9}=3,\sqrt{1}=1$,$3+1=4$)。代入同一个目标函数:$9/0.75+1/0.25 = 12+4=16$——严格更低,16 < 20,证实平方根规则确实比看似更”直觉”的线性比例规则表现更好,尽管线性分配看起来是让重要时间步更容易被纳入的更”显然”的做法。直觉上的原因:目标函数 $w_t/\xi_t$ 对小 $\xi_t$ 的惩罚是凸的、加速增长的(像 $1/\xi_t$ 一样爆炸,而不是线性增长),所以把 $\xi_t$ 从一个已经很小的值移开的边际收益,比把一个已经很大的 $\xi_t$ 再抬高的边际收益要大——平方根分配恰好找到了让边际 $w_t/\xi_t^2$ 惩罚(来自稳定性条件 $\partial/\partial\xi_t=-w_t/\xi_t^2+\lambda=0$)在所有时间步之间完全均衡的那个平衡点,这正是拉格朗日推导所形式化、也被这个数值检验具体证实的东西。

为什么这套凸优化/单调性机制值得费这个功夫(设计选择讨论)

值得暂停一下,问一个问题:为什么不跳过所有这些优化,直接用一个简单的固定衰减策略 $\xi_t = \exp(-\lambda(t-1))$?论文实际上把这个方案当作一个基线来测试——“PS-PPO (Time-Prior)“——论文的表 4/图 3 显示它的表现并不比均匀随机截断更好。原因,由上面的推导说得很清楚:固定衰减策略假定了”一段轨迹变得可预测的速度”这个”形状”对每个 prompt 都是一样的,但图 1 本身就说明这是错的,即便只在两个基准之间比较(AIME 2024 在长度的 25% 处就已经稳定,MATH-500 是 38%)——更不用说同一个基准内不同题目之间的差异了。上面推导出的、依赖 prompt 的最优截断分布 $\xi_{1:T}(x)$,恰恰不是一个只依赖 $t$ 的固定函数;它通过 $w_t^\theta(x)$ 依赖于 $x$,而 $w_t^\theta(x)$ 是从当前 batch 里那个具体 prompt 的 rollout 重新估计出来的。这就是”如果不这样做会怎样”问题的答案,回答了为什么更重的优化机制值得投入:一个与上下文无关的启发式规则,无法适应”一段轨迹在何时变得可预测”本身是这道具体题目的属性,而不是轨迹长度的一个普适常数这一事实。

一个真正需要合并的 PAV 例子

前面给出的数值例子(见后文”从权重到采样截断点的完整数值例子”一节)刻意构造成原始分数 $\sqrt{w_t}$ 天生就已经是单调的,所以 PAV 无用武之地。既然 PAV 是这篇论文相对于标准拉格朗日松弛真正引入的唯一新算法组件,值得走一遍单调性确实被违反、必须修复的例子。

假设 $T=4$,原始权重 $w_1=1$,$w_2=9$,$w_3=4$,$w_4=1$——注意 $w_2 > w_1$,所以无约束的分数 $\sqrt{w_t} = 1, 3, 2, 1$ 不是非增的(从 $t=1$ 到 $t=2$ 反而升高了,然后才降下来),如果直接拿这个当 $\xi_t$ 用,就会违反生存函数所需的单调性。

PAV 第一步——初始化单点块。 每个时间步从自己独立的一个块开始:块分数是 $[1, 3, 2, 1]$,块长度都是 1,块权重就是原始 $w_t$:$[1, 9, 4, 1]$。

PAV 第二步——扫描违规。 从左到右比较相邻块:块 1(分数 1)对比块 2(分数 3)——违规,因为需要非增顺序但 $1 < 3$。合并块 1 和块 2:新块长度 $L=1+1=2$,权重 $W=1+9=10$,分数 $s=\sqrt{W/L}=\sqrt{10/2}=\sqrt{5}\approx 2.236$。

PAV 第三步——合并后重新检查。 现在块是 $[\{1,2\}: \text{分数}=2.236, L=2, W=10]$,$[\{3\}: \text{分数}=2, L=1, W=4]$,$[\{4\}: \text{分数}=1, L=1, W=1]$。检查块 1(2.236)对块 2(2.0):$2.236 \geq 2.0$——没问题,不违规。检查块 2(2.0)对块 3(1.0):$2.0\geq 1.0$——没问题。块分数序列 $[2.236, 2.0, 1.0]$ 现在是非增的了,PAV 终止。

PAV 第四步——赋值最终的 $\xi_t^\star$。 取预算 $B=2$(举例),归一化分母是 $\sum_m L_m s_m = 2(2.236) + 1(2.0) + 1(1.0) = 4.472+2.0+1.0 = 7.472$。于是:$\xi_1^\star = \xi_2^\star = B\cdot 2.236/7.472 = 0.599$(这两个时间步得到相同的 $\xi$,因为 PAV 强制它们被合并进同一个块),$\xi_3^\star = B\cdot 2.0/7.472 = 0.535$,$\xi_4^\star = B\cdot 1.0/7.472=0.268$。

检查一遍。 单调性:$0.599 \geq 0.599 \geq 0.535 \geq 0.268$——非增(合并块内部相等),满足生存函数的要求。预算:$0.599+0.599+0.535+0.268 = 2.001\approx B=2$——在舍入误差范围内正确。

值得停下来体会的解读: 时间步 1,单独看原始权重最小($w_1=1$),最后却获得了与原始权重最大的时间步 2($w_2=9$)完全相同的纳入概率——因为一旦检测到违规,单调性约束就把它们强制归入了同一个块。这就是施加单调性约束的真实代价:某些个别时间步会与邻居”合并”,失去获得量身定制的纳入概率的能力,合并的严重程度恰好与原始重要性信号违反所需非增形状的程度成正比。这也正是为什么论文图 1 的动机(重要性在轨迹推进过程中总体上大致递减)对 PAV 在实践中不至于严重损害解的质量很重要——如果真实训练中原始重要性高度非单调,PAV 就会把轨迹的大部分合并在一起,稀释掉优化本来想要争取的精细分辨率。

让权重可计算,而不需要反向传播

只用前向传播的 Score-Norm 代理量

上面的设计存在一个”鸡生蛋、蛋生鸡”的问题:计算 $w_t^\theta(x) = \mathbb{E}[\hat A_t^2\|\nabla_\theta\log\pi_\theta(o_t\mid s_t)\|^2\mid x]$ 需要对全部 $\theta$ 求出完整的梯度范数 $\|\nabla_\theta\log\pi_\theta(o_t\mid s_t)\|$——这恰好需要 PS-PPO 想要避免的完整反向传播,而且是对每一个候选截断决策都要算一次。论文的解决方案是,只用输出层对梯度的贡献来近似,这可以只从前向传播的量算出来。

闭式输出层 score 的推导(附录 D)。 对标准的 softmax 输出层,$z_t = Wh_t + b$,$p_t = \mathrm{softmax}(z_t)$,其中 $h_t$ 是最后一层的隐藏状态(前向传播已经算出来了)。对实际采样出的 token $o_t$,它对数概率对输出层权重矩阵 $W$ 的梯度有一个众所周知的闭式:

$$\nabla_W\log p_t(o_t) = (e_{o_t}-p_t)\,h_t^\top, \qquad \nabla_b\log p_t(o_t) = e_{o_t}-p_t,$$

其中 $e_{o_t}$ 是采样 token 的独热向量。这就是标准的 softmax 交叉熵梯度恒等式:关于 logits 的梯度恰好是”独热向量减去预测分布”,完全不需要反向传播就能推导出来,因为它是前向传播已经产生的量($p_t$)和采样 token 身份($e_{o_t}$)的闭式函数。取 Frobenius 范数的平方:

$$\|\nabla_W\log p_t(o_t)\|_F^2 = \|h_t\|_2^2\,\|e_{o_t}-p_t\|_2^2, \qquad \|e_{o_t}-p_t\|_2^2 = 1-2p_t(o_t)+\|p_t\|_2^2.$$

第二个等式展开了 $\|e_{o_t}-p_t\|_2^2 = \|e_{o_t}\|_2^2 - 2\langle e_{o_t},p_t\rangle + \|p_t\|_2^2 = 1 - 2p_t(o_t) + \|p_t\|_2^2$,利用了 $e_{o_t}$ 是独热向量($\|e_{o_t}\|^2=1$)以及 $\langle e_{o_t},p_t\rangle = p_t(o_t)$(取出采样 token 的概率)。合并起来就得到论文的代理量:

$$\gamma_t(s_t,o_t) := \|h_t\|_2^2\big(1-2p_t(o_t)+\|p_t\|_2^2\big).$$

高效处理看似昂贵的 $\|p_t\|_2^2$ 项。 朴素地计算 $\|p_t\|_2^2 = \sum_i p_{t,i}^2$ 需要把整个(10 万以上词表大小的)softmax 分布完整算出来并逐项求平方,这不是不可承受的开销,但附录 D 给出了一个更巧妙的 log-sum-exp 技巧可以避免这一点:定义 $\ell_1 = \log\sum_i e^{z_{t,i}}$(标准的对数配分函数,计算 token 对数概率时已经算过)和 $\ell_2 = \log\sum_i e^{2z_{t,i}}$(对双倍 logits 再做一次 log-sum-exp),那么 $\|p_t\|_2^2 = \exp(\ell_2 - 2\ell_1)$——这直接来自 $\sum_i p_{t,i}^2 = \sum_i\left(\frac{e^{z_{t,i}}}{\sum_j e^{z_{t,j}}}\right)^2 = \frac{\sum_i e^{2z_{t,i}}}{(\sum_j e^{z_{t,j}})^2} = \exp(\ell_2-2\ell_1)$。$\ell_1$ 和 $\ell_2$ 都是对 logits 的廉价标量归约,可以在已经要计算 token 对数概率的同一次前向传播中一并算出——不需要反向传播,也不需要在整个词表上做显式的逐元素平方求和循环(除了一次 log-sum-exp 的变体)。

为什么这个代理量是站得住脚的,而不只是图方便。 由于 $(W,b)$ 是全部可训练参数 $\theta$ 的一个真子集,而平方范数在独立参数块之间可以加性分解,$\|\nabla_\theta\log\pi_\theta(o_t\mid s_t)\|^2 \geq \|\nabla_{W,b}\log p_t(o_t)\|^2$——输出层的范数永远是真实全参数 score norm 的一个下界,而不是一个不相关的相关量。论文进一步用实证验证了这个下界与真实范数贴合得很紧,而不是松散地相关:

双对数图上 0.9968 的皮尔逊相关系数是一个很强的实证验证——它说明尽管输出层范数在理论上只是一个下界、并不精确匹配,但”哪些 token 梯度更大”这种相对排序几乎被完美保留了下来,而这正是截断设计优化实际需要的性质(它只需要正确地对时间步的重要性排序,不需要匹配它们绝对的梯度大小)。

奖励不确定性项及其上界推导

$w_t^\theta(x)$ 需要的另一个原料,是一旦优势平方项 $\hat A_t^2$ 通过全期望公式和一个均场近似进入之后,对 $\mathrm{Var}(R\mid s_t)$——即”给定前缀状态 $s_t$,最终奖励还有多不确定”——的一个估计。对于二值奖励来说,这是一个 Bernoulli 变量的方差,论文推导出了一个可计算的上界(附录 E),而不是尝试直接估计方差(那需要对每个前缀做很多次后缀 rollout——这正是图 1 验证实验专门用的、代价高昂的做法,不是每步训练都负担得起的)。

第一步——Bernoulli 方差的一个通用上界。 对 $R\in\{0,1\}$,$\Pr(R=1\mid s_t)=p_t(s_t)$,方差是 $p_t(1-p_t)$,一个标准事实是 $p_t(1-p_t)\leq\min\{p_t,1-p_t\}$(方差不超过两者中较小的那个——这是真的,因为当 $p\leq 1-p$ 时,$p(1-p)-\min(p,1-p) = p(1-p) - p = -p^2 \leq 0$,反过来对称)。用恒等式 $\min\{a,b\}=\frac{1}{2}(a+b-|a-b|)$,取 $a=p_t,b=1-p_t$,得到 $\min\{p_t,1-p_t\} = \frac{1}{2}-\frac{1}{2}|2p_t-1|$。

第二步——对随机前缀状态取期望,把绝对值项重写成一个类似全变差距离的量。 对两边取 $\mathbb{E}_{S_t\mid x}$,并利用 $2p_t(s_t)-1=\Pr(R{=}1\mid s_t)-\Pr(R{=}0\mid s_t)$ 展开 $|2p_t(S_t)-1|$,经过一个简短的代数重写(论文附录 E 的公式 20),把期望转化成了(前缀状态, 奖励)这个联合分布上的一个 $\ell_1$ 距离:

$$\mathbb{E}[|2p_t(S_t)-1|\mid x] = \sum_{s_t}\big|\Pr(s_t,R{=}1\mid x)-\Pr(s_t,R{=}0\mid x)\big|.$$

第三步——把前缀状态映射到下一个 token 的分布,让这个量变得可估计。 前缀状态 $s_t$ 上的分布在训练中并不方便直接估计(它生活在一个极大的状态空间里),所以论文用三角不等式把这个量进一步限定到用下一个 token 在旧策略下的分布来表达——这在 rollout batch 里已经有一个自然的经验估计量:

$$u_t(x) := \frac{1}{2}-\frac{1}{2}\big\|p(x)\bar\pi_G(\cdot\mid t,x)-(1-p(x))\bar\pi_B(\cdot\mid t,x)\big\|_1,$$

其中 $p(x)=\Pr(R{=}1\mid x)$ 是该 prompt 的经验成功率(从 $K$ 段回复的组里估计),$\bar\pi_G,\bar\pi_B$ 是旧策略下按状态平均的下一个 token 分布,分别在组内成功和失败的 rollout 上计算。最终的保证 $\mathbb{E}[\mathrm{Var}(R\mid S_t)\mid x]\leq u_t(x)$ 是一个真正的数学上界(不是一个方向未知的启发式近似)——上面这一连串不等式(Bernoulli 方差上界 $\to$ 三角不等式 $\to$ 用下一个 token 分布重写)全程都是单调的,所以 $u_t(x)$ 只可能高估、绝不会低估真实的奖励不确定性。实践中,$p(x)$、$\bar\pi_G$、$\bar\pi_B$ 都可以从组内优势估计已经用到的同一批 $K$ 段 rollout 中经验估计出来——不需要额外的 rollout,不需要辅助的奖励不确定性模型。

拼在一起:最终优化问题与完整算法

把 score-norm 代理量 $\gamma_t(x,t)$(在 batch 内聚合为 $\bar\gamma_t(x,t):=\mathbb{E}_{s_t\mid x,\,o_t\sim\pi_{\theta_{\text{old}}}(\cdot\mid s_t)}[\gamma_t(s_t,o_t)]$)和奖励不确定性上界 $u_t(x)$ 结合起来,论文把最终的逐时间步设计权重定义为两者的乘积 $w_t(x) := \bar\gamma_t(x,t)\,u_t(x)$,并求解:

$$\arg\min_{\xi_{1:T}}\ \sum_{t=1}^{T}\frac{\bar\gamma_t(x,t)\,u_t(x)}{\xi_t} \quad \text{s.t.}\quad \sum_{t=1}^{T}t(\xi_t-\xi_{t+1})=B,\ \ \xi_{T+1}:=0,\ \ 1\geq\xi_1\geq\cdots\geq\xi_T>0.$$

按照上面走过的推导,这个问题由拉格朗日闭式 $\xi_t \propto \sqrt{\bar\gamma_t(x,t)u_t(x)}$ 加上 PAV 恢复单调性求解。

算法 1(Prefix-Sampling PPO),逐行复现并加注:

需要: 策略 pi_theta; prompt 分布 p_Q; 奖励函数 R;
      最大训练步数 N; 旧策略刷新周期 F;
      组大小 K; 最大回复长度 T; 算力预算 B。

 1: 初始化 theta, 并设 theta_old <- theta。
 2: for 迭代 n = 1 to N do
 3:     采样一个 prompt x ~ p_Q。
              # 与任意无 critic 的 RLHF 方法相同
 4:     从 pi_theta_old(. | x) 采样 K 段回复 {tau^(i)}。
              # 相同 -- 这是 PS-PPO 不去动的 rollout 阶段
 5:     计算终局奖励 {R^(i)}。
              # 相同 -- 奖励函数不受影响
 6:     用组内基线计算无 critic 优势 {A_hat^(i)}。
              # 相同 -- 例如 GRPO 风格的组内相对优势
 7:     从 {(tau^(i), R^(i))} 计算逐时间步代理量 {u_t(x), gamma_bar_t(x,t)}。
              # 新增: 只用前向传播的 score-norm 代理量 (公式 6) + 奖励不确定性上界 (公式 7),
              #      两者都从第 4 步已经采样好的同一批 K 段 rollout 里估计 -- 不需要额外数据
 8:     设截断权重 w_t(x) <- gamma_bar_t(x,t) * u_t(x), t = 1..T。
              # 新增: 把两个代理量合并成设计权重
 9:     求解预算约束下的设计问题 (公式 8), 计算单调截断概率 xi_1:T。
              # 新增: 闭式 sqrt(w_t) 分配 + PAV 恢复单调性 (见上面的伪代码)
10:     从 xi_1:T 诱导出的截断分布中采样截断点 {H^(i)}。
              # 新增: 每段回复一次随机抽取
11:     在更新阶段, 只对 t <= H^(i) 的 token 做反向传播,
              并把 token 损失按 1/xi_t 重加权。
              # 新增: 这才是真正省算力的地方 -- 更短的反向传播计算图
12:     用截断且重加权的梯度估计量,按 PPO 更新 theta (论文中的公式 9)。
              # 相同的 PPO 更新机制, 只是作用在更少的 token 上
13:     if n mod F == 0 then
14:         theta_old <- theta
              # 相同 -- 标准的旧策略刷新周期
15:     end if
16: end for
17: return 优化后的策略 pi_theta。

把这些注释放在一起看,这个算法对已有无 critic RLHF 流水线的实际改动,集中在第 7 到第 11 步——rollout 生成、奖励打分、优势计算、以及 PPO 更新规则本身都不需要改动。这正是论文能把下面的效率数字说成”几乎即插即用的替换”、而不是需要重新设计整套训练系统的原因。

为什么截断前向传播也能省显存,不只是省时间

有必要明确说清楚表 1 里提到的 15%-17% 峰值显存下降背后的机制,因为”更短的反向传播”和”更少的显存”是相关但不完全相同的说法,混为一谈会漏掉一部分图景。

在标准 Transformer 前向传播中,每一层的中间激活值(注意力分数、前馈隐藏状态、层归一化后的输出等)都必须保留在显存里,直到反向传播用到它们来计算梯度为止——这就是众所周知的激活显存瓶颈,梯度检查点等技术正是针对这一点。存储这些激活值的显存开销大致随序列长度线性增长:长度为 $T$ 的序列,需要存储的激活显存大约是长度为 1 的序列的 $T$ 倍(每层、每个 batch 元素),因为注意力机制以及后续每个 token 的计算都必须为全部 $T$ 个位置保留下来。

当 PS-PPO 在前向传播之前就把一段回复截断到采样出的截断点 $H^{(k)} < T$(而不是先跑完全部 $T$ 个 token 的前向传播、事后再丢弃一部分),所需的激活显存就正比于 $H^{(k)}$,而不是 $T$。对一批截断点各不相同(随机采样)的回复求平均,期望激活显存正比于 $\mathbb{E}[H] = B$,而不是(通常大得多的)$T_{max}$——这正是图 2(c)峰值 GPU 显存曲线里,PS-PPO 相比全序列基线在整个训练过程中都保持更低显存的直接机制解释,不只是在训练末期才体现。

这也正是为什么只做掩码的基线(S-GRPO、带 forking-token 的 DAPO)没有表现出类似的显存下降:事后把损失掩码掉,并不能减少全长前向传播过程中已经必须分配和保留的激活显存——等到掩码被应用的时候,显存早就已经花出去了。这与解释时间节省(表 1)和显存节省(图 2(c))的底层机制其实是同一个(截断计算图本身,而不是事后掩盖一个全长计算图)——它们是同一个根本原因的两个症状,而不是两个独立的设计成果。

实验:数字实际说明了什么

效率:时间和显存节省从哪里来

面板(d)是让论文核心效率主张一目了然的那张图:对 S-GRPO 和带 forking-token 的 DAPO 来说,损失非零的 token 数量(浅色柱)明显小于实际被反向传播的 token 数量(深色柱)——这些方法只是掩盖了哪些 token 贡献梯度,但前向和反向传播仍然要跑过完整的、未被掩码的序列长度,才能先算出那些事后被掩成零的损失值。PS-PPO 的柱子是唯一一对损失 token 和反向传播 token 一起缩小的,因为截断发生在前向传播开始之前(喂给网络的序列本身就更短),而不是对一个已经算完的损失张量做事后掩码。这正是第 15 条要求的”为什么是这个设计选择,显然的替代方案会怎么做”的讨论:PS-PPO 的显然替代方案正是 S-GRPO 和带 forking-token 的 DAPO 已经在做的事(先算完一切,再把不想要的贡献清零),而这个替代方案没能省下多少墙钟时间的原因是,现代加速器硬件不管前向/反向矩阵乘法的输出之后是否被损失掩码清零,都要为这些计算买单——已经花掉的 FLOPs 是没法收回的。

把表 1 逐项拆开看:PS-PPO 的前向开销(0.32 秒)已经低于每一个基线的前向开销(0.82-1.12 秒),这是合理的,因为 PS-PPO 的序列在前向传播开始之前就已经更短了——截断不只是一个反向传播阶段的优化。反向开销(1.01 秒)同样是四者中最低的。PS-PPO 唯一多花钱的地方是”计算 $\xi_{1:T}$“这一列(0.43 秒,而完全不做 prompt 相关截断设计的方法这一项是 N/A)——论文诚实地把这笔开销计入总时间,而不是把它藏起来,说明即便算上这笔开销,净收益依然明显。

准确率:效率是否是以牺牲准确率为代价换来的?

这里最关键的比较是”PS-PPO(Optimized, B=128)“与最强基线之间的对比。在 Llama-3.1-8B-Instruct 上,PS-PPO(Optimized)在 MATH500 上取得单项最佳成绩(47.6,超过带 forking-token 的 DAPO 的 47.0),在 AMC23 上(32.5,与 DAPO 和带 forking-token 的 DAPO 打平),在其余四个基准上具有竞争力(不是最好,也不是最差)。在 Qwen2.5-Math-7B 上,PS-PPO(Optimized)取得 AIME25 最佳成绩(13.3,而次好的方法只有 10.0),同时在其他基准上也保持竞争力。同样有信息量的是 PS-PPO 家族内部的对比:“PS-PPO(Uniform)”、“PS-PPO(Time-Prior)”、“PS-PPO(Heuristic)“在几乎每个基准和两个基座模型上,都以一个一致的、不小的差距落后于”PS-PPO(Optimized)“——这说明收益恰恰可以归因于求解了这个凸设计问题,而不仅仅是”以某种方式截断了训练”。一个持怀疑态度的读者自然会问的问题——“也许任何合理的截断方案效果都差不多,这一整套优化机制是多余的”——正是被这个内部消融实验直接回答了,数据给出的答案是否定的。

随回复长度的可扩展性

随着思维链训练朝更长回复的方向发展,论文直接测试了 PS-PPO 的优势会随回复长度增长还是缩小,扫过 $T_{max}\in\{1024,2048,4096\}$:

$T_{max}$	方法	每步时间(秒)	平均准确率(MATH500, AIME24, AIME25)
1024	PS-PPO	1.77 ± 0.02	38.3
1024	S-GRPO	2.66 ± 0.01	37.1
1024	DAPO	3.23 ± 0.01	37.3
4096	PS-PPO	2.39 ± 0.04	42.1
4096	S-GRPO	6.70 ± 0.04	39.7
4096	DAPO	7.78 ± 0.04	42.2

图 3(论文表 3,改写成对比表):在 $T_{max}=4096$ 时,PS-PPO 每步时间(2.39 秒)比 S-GRPO(6.70 秒)快 2.8 倍,比 DAPO(7.78 秒)快 3.3 倍,同时平均准确率(42.1)与 DAPO(42.2)在噪声范围内没有差异,明显领先 S-GRPO(39.7)。

这里与设计相关的观察是,随着 $T_{max}$ 增大 3 倍,PS-PPO 每步时间增长得多么少(1.77 秒 $\to$ 2.39 秒,增长 35%),相比之下基线每步时间增长了多少(DAPO:3.23 秒 $\to$ 7.78 秒,增长 141%)——这是 PS-PPO 期望反向传播长度由预算 $B$ 而非 $T_{max}$ 本身决定这一机制的直接后果:随着回复变长,PS-PPO 只是截断了每段回复中更大的比例,而基线的反向传播开销大致随完整回复长度线性增长,不管结果在中途变得有多可预测。

在匹配预算下,单独考察截断策略本身

为了回答”收益是来自每步算力更少,还是来自对同样算力的更聪明分配”这个具体问题,论文在匹配的期望反向传播长度下做了受控对比($B=512$,$T=1024$,所以均匀截断按构造给出 $\mathbb{E}[H]=T/2=512$,与其他策略的预算精确匹配):

这一对图直接回答了截断设计这一选择的”显然的替代方案会怎么做”这个问题:如果固定花费的算力数量(相同的期望反向传播 $B$ 个 token),只改变这些算力花在哪里(哪些时间步被纳入),Optimized/Heuristic 策略在墙钟时间上到达同样最终奖励的速度明显快于 Uniform 或 Time-Prior——也就是说,收益不仅仅是”PS-PPO 花的算力更少”(那是表 1/图 2 比较的内容),还在于”在固定算力预算下,PS-PPO 把这个预算在各时间步之间的具体分配方式,可以证明比朴素的替代方案更好”(这里比较的内容)。论文特别指出了 Time-Prior 的失败之处:仅仅偏爱较早的时间步(而不依据实际的奖励不确定性信号做条件化)的表现并不比纯 Uniform 更好——这证实了真正起作用的是设计原则本身(按 $w_t(x)$ 分配概率质量,而不是按原始位置 $t$),而不是某种泛泛的”偏爱早期 token”的近因偏差。

超参数敏感性:预算 $B$ 与组大小 $K$

论文报告了两项值得简要总结的消融实验(论文中的表 4 和表 5):扫描预算 $B\in\{64,128,256,512\}$,准确率(平均)从 $B=64$ 时的 34.0 上升到 $B=128$ 时的 37.9,之后基本进入平台期($B=256$ 时 37.5,$B=512$ 时 37.8),而每步训练时间持续线性攀升(1.15 秒 $\to$ 1.77 秒 $\to$ 2.01 秒 $\to$ 2.56 秒)——这清楚地标出了在 $B=128$ 附近的收益递减拐点,论文也把这个值作为默认设置。扫描组大小 $K\in\{2,4,8,16\}$ 显示出一个相似但不同的模式:准确率一路提升到 $K=16$(34.2 $\to$ 38.5),但论文认为 $K=8$ 是实际使用中的甜蜜点,因为它以远低于 $K=16$ 的每步训练成本(1.77 秒对 2.24 秒)取得了接近 $K=16$ 的表现(37.9 对 38.5)——还有一个更微妙、值得明确指出的点:$K$ 直接影响 $u_t(x)$ 估计本身的质量——当 $K$ 较小时,用来估计 $\bar\pi_G,\bar\pi_B$ 的成功/失败拆分经常一侧样本太少(或者所有 rollout 都落在同一侧,变成一个完全”退化”的、不管什么截断策略都没有学习信号的 batch),让 $u_t(x)$ 变得嘈杂,进而把截断分布推向一个信息量不足的近似均匀形状——这是 $K$ 与截断设计机制之间的一个二阶交互作用,如果只看头条准确率数字很容易被忽略。

论文承认的局限性与边界条件

• 所有主要实验都使用二值正确性奖励。 连续奖励的 RLHF 实验(附录 G,HH-RLHF 和使用学习到的奖励模型的 IMDB 正向生成任务)需要对奖励做一个临时的、基于 sigmoid 的二值化 $p=\sigma(r)$,纯粹是为了接入本来是为二值场景推导出来的 $u_t(x)$ 估计器。论文明确标注了这是把一个估计器用在了它原始推导假设之外的场景。
• 对角线/迹方差代理量丢弃了跨时间步的协方差。 论文坦率地承认,精确的方差展开式中包含来自自回归依赖关系的跨时间步协方差项,保留这些项会让 $\xi_t$ 在时间步之间耦合起来,破坏掉闭式可解性——但这意味着”最优的” $\xi_{1:T}$ 只是相对于可处理的代理量而言的最优,而不是相对于真实方差而言的最优。
• 只用前向传播的 score-norm 代理量,是真实梯度范数的一个严格下界,只在 Qwen2.5-Math-7B 上的相关性研究(图 4)中做过验证;论文没有证明这个代理量在输出层设计明显不同的架构上(例如绑定嵌入、最后一层做混合专家路由)或不同规模下,能保持同等的保真度。
• 奖励不确定性上界 $u_t(x)$ 是一个套了两层不等式的松弛(先是 Bernoulli 方差上界,再是三角不等式意义上的状态到下一 token 映射)——论文证明了它是一个有效上界,但没有刻画它在实践中通常有多松,只说明它可以廉价地估计出来。
• 所有主要数学推理实验都用同一个数据集家族(MATH,在 MATH500/AMC23/CollegeMath/MinervaMath/AIME24/AIME25 上评测)和两个模型规模(7-8B)。 是否能推广到明显更大的模型(70B+),或者推广到 prefix 可预测性可能表现不同的非数学领域(比如开放式创意写作、有延迟奖励的多轮对话),没有经过实证测试。
• 迟来的”顿悟”被承认但没有被消除。 附录 A 自己的分析显示,即便到了 75% 的 prefix 位置,仍有 7%-9% 最终正确的回复,其 prefix 条件成功率仍低于 0.8——该方法用随机(而非确定性)截断,正是为了对冲这种情况,通过仍以非零概率允许更晚的截断点来实现;但论文没有报告,如果换成一个更激进、事实上消除了这种对冲的预算,会损失多少准确率。

批判性分析:不足与可改进之处

不够有说服力的证据。 首先,“33%-45% 训练时间下降”和”15%-17% 显存下降”这两个头条数字(表 1、图 2)是在排除 rollout/生成时间的前提下测出来的——论文明确说明了这一点,但这意味着报告的加速只适用于梯度更新阶段,而不是整个 RLHF 训练流程的墙钟成本,而在实际中,rollout/生成(用推理引擎采样回复)往往是整体成本里占大头的部分,对长思维链回复尤其如此。论文从未报告一个包含 rollout 时间的端到端墙钟对比,而这样的对比会给出一个更保守(也对实践者更相关)的实际加速图景。其次,表 2 的准确率对比显示,PS-PPO(Optimized)只在 12 个”基准×基座模型”格子中大约一半上获胜或打平,在其余格子上略微落后于 DAPO 或带 forking-token 的 DAPO(例如:CollegeMath 上 Llama 为 33.5 对 DAPO 的 34.2;Qwen 上 23.9 对带 forking-token 的 DAPO 的 25.0)——论文把这称为”可比的准确率”,鉴于报告的标准误没有按格子展示(主表里只说明”三次独立运行取平均”,没有给出置信区间),这个说法是可以站得住脚的,但没有逐格方差条,读者根本无法判断这些微小的差距是否具有统计意义,还是纯粹是噪声,这是证据链条里一个真实的缺口,考虑到很多数字确实非常接近。第三,没有任何消融实验单独隔离出只用前向传播的 score-norm 代理量(公式 6)相对于直接使用真实score norm(这需要昂贵的反向传播,但至少可以在一个较小模型上跑一次,作为受控对比来检验)的边际贡献——图 4 的高相关性对排序保真度来说是令人安心的,但论文从未闭环地展示出:用精确的 score norm(接受额外开销,纯粹作为诊断)训练,会得到与用代理量训练相同的截断分布和相同的下游准确率。

论文低估或未充分探讨的局限。 论文把 Time-Prior 基线的失败(“单纯偏爱更早的时间步是不够的”)当作证据,说明真正起作用的是那个经过学习的、依赖 prompt 的信号——但它没有测试一个能加强或削弱这一论断的中间假设:一个固定的、数据集层面(非 prompt 条件)的经验 prefix 可预测性曲线,用图 1 本身的方法在一批留出的训练 prompt 上估计一次,然后在推理时统一应用,不做任何逐 prompt 或逐 batch 的重新估计。如果这样一个”数据集平均、但非平凡形状”的截断计划表现能与完全 prompt 条件化的 Optimized 变体相当,那就说明 prompt 条件化本身(需要每一步都重新估计 $u_t(x)$ 和 $\bar\gamma_t(x,t)$ 的那部分昂贵计算)对论文论述的贡献,可能不如论文所暗示的那么大,而相对于 Time-Prior 的大部分收益,可能仅仅来自使用了正确的非单调形状(匹配图 1 实际的 S 型曲线),而不是来自逐 prompt 的适应。论文的消融设计(Uniform 对 Time-Prior 对 Heuristic 对 Optimized)无法区分这两种假设,因为 Time-Prior 的失败既可以解释为”形状错了”,也可以解释为”没有逐 prompt 适应”,论文从未把它们隔离开来。另外,论文的效率主张是在 7-8B 模型、8 张 A100 上验证的——没有讨论计算 $\xi_{1:T}$ 的相对开销(论文报告在这个规模下大致固定在每步约 0.43 秒)在明显更大的模型或 batch 规模下会如何变化,而每步的前向/反向成本会随规模增长,但 $\xi$ 计算成本(受词表大小操作约束,而非全模型规模操作约束)可能不会成比例增长——这可能意味着相对效率优势在更大规模下要么增大要么缩小,论文对此没有给出任何讨论。

具体的改进建议。 (1)至少针对一个代表性配置,报告一个包含 rollout/生成时间的端到端墙钟对比,让实践者能判断一次完整训练运行的实际加速比,而不仅仅是孤立的更新阶段数字。(2)在表 2 中加入逐格的置信区间或标准误,鉴于很多准确率对比数字非常接近——三次独立运行是计算标准误的合理次数,报告出来会大大加强”准确率可比”这一论断。(3)运行上面建议的受控消融——用精确的(反向传播计算出的)score norm 对比只用前向传播的代理量训练一个较小模型,固定其他一切不变——直接量化代理量在最终准确率上到底(如果有的话)付出了多少代价,补上图 4 相关性研究留下的缺口。(4)加入一个”固定非单调计划表”基线(图 1 曲线的数据集平均版本,不做逐 prompt 重新估计),干净地把 Optimized 变体相对于 Time-Prior 的优势中,“形状正确”与”逐 prompt 适应性”各自的贡献区分开来。(5)把回复长度可扩展性研究(目前 $T_{max}\in\{1024,2048,4096\}$,表 3)扩展到明显更长的回复(8k-32k token),因为这正是——很长的思维链推理轨迹——论文自己的动机(图 1 的 prefix 可预测性论证)所暗示该方法应该拥有最大相对优势的区间,但目前尚未经过测试。

何时最该用 PS-PPO?一张速查表

下面是一张面向决策的简要总结,帮你判断这篇论文引入的机制在什么场景下最值得投入实现成本,什么场景下可能不值得:

场景	PS-PPO 的相关性
长思维链强化学习训练(每段回复上千个 token)	高——表 3 显示,相对于只做掩码的基线,加速比从 $T_{max}=1024$ 时的约 1.5 倍,提升到 $T_{max}=4096$ 时的 2.8-3.3 倍,而且论文自己的动机说明这个趋势在更长长度下应该会继续
二值/可验证奖励(数学、代码正确性、精确匹配类任务)	高——这正是奖励不确定性代理量 $u_t(x)$ 原生推导所针对的场景
基于连续奖励模型的 RLHF(通用的有用性/无害性)	中等——按附录 G 可行,但依赖在原生推导之上再叠加一层 sigmoid 二值化近似
Rollout/生成阶段主导总训练墙钟时间(例如推理引擎很慢、模型小但采样开销大)	较低——PS-PPO 只优化梯度更新阶段,不涉及 rollout;如果 rollout 本来就是瓶颈,相对收益会缩小
短回复(只有几十个 token,例如单句分类式奖励)	低——一开始就没有多少冗余的尾巴可截,因此该方法的核心动机(图 1 的 prefix 可预测性曲线)就不那么适用了
已经在使用只做掩码的 token 选择方法(S-GRPO、带 forking-token 的 DAPO)	高——图 2(d) 显示这些方法减少了 loss token 但没有减少反向传播 token;改用 PS-PPO 真正的截断方法,很可能在不牺牲已经获得的 token 选择收益的前提下,进一步实现额外节省
使用带学习到的 token 级 critic 的 actor-critic 方法(非无 critic)	未测试——论文的具体推导假设了广播式(在 $t$ 上保持不变)的优势;适配到随 $t$ 变化的优势上,上文已经标注为一个自然但尚未探索的延伸方向

作者在 github.com/doohwan383/PS-PPO 关联了一份实现,基于 Hugging Face 的 Open-R1 代码库扩展而成。既然本文通篇都在强调 PS-PPO 对已有训练循环的实际改动范围很小、且相互隔离(两个只用前向传播的代理量、基于 PAV 的截断求解器,以及更新阶段的截断/重加权逻辑——见上文算法 1 的注释),验证本文任何一个具体推导最直接的方式,就是去那份代码库里找到对应的函数,直接对照论文的公式检查,而不需要从头复现一次完整的训练流程。凸设计问题本身(公式 8,通过平方根分配加 PAV 求解)是一个小巧、自包含的数值程序,可以独立地用上面给出的手算例子做单元测试,不依赖于任何实际的语言模型训练。

结语

PS-PPO 把一个真正有用的观察变成了可操作的方法:如果一段推理轨迹最终的奖励往往在中途就已经大体确定,那么无 critic 的 RLHF 方法就没有理由必须为整段轨迹的每一次梯度更新都支付完整的前向/反向算力——这只是朴素的广播式优势表述碰巧隐含的做法,除非有人搭建出安全截断所需要的机制。论文真正的贡献正是这套机制:一个按纳入概率重加权、对任意合法截断分布都可证明无偏的估计量;一个在算力预算约束下、最小化方差代理量来选择这个分布的凸设计问题;一个让计算设计权重本身变得廉价的只用前向传播的代理量;以及一个不需要额外 rollout 就能估计的奖励不确定性上界。端到端地测量,这转化成了更新阶段每步训练时间实实在在减少 33%-45%,峰值显存减少 15%-17%,准确率在统计意义上与强无 critic 基线接近(在部分基准上甚至更好)——而且回复越长,这个优势就越明显,这正是长思维链强化学习训练正在走向的方向。本文批判性分析部分提出的悬而未决的问题——包含 rollout 的端到端墙钟对比、逐格的统计显著性、把”形状正确”和”逐 prompt 适应性”区分开——是任何想直接在这项工作基础上继续构建的人自然而然的下一步实验,而不是怀疑核心机制本身的理由,因为这个机制在每一个非平凡的步骤上都经过了仔细的推导和验证,而不只是被断言出来。

对于正在纠结是否要投入工程时间去采用这个具体方法、还是干脆等待领域在同一个想法上的下一轮迭代的读者来说,这篇论文最持久的部分,与其说是具体的代理量或具体的 PAV 求解器,不如说是它建立起来的模式:每当一个 RL 训练信号被同等地广播到许多计算单元上时(这里是 token,但这个模式也能推广到其他有类似冗余的逐单元计算场景),都值得明确地检查一下,是否存在一个无偏的、受算力预算约束的、方差最小化的采样方案,能够以一小部分完整计算的成本恢复大部分完整计算带来的信息收益——如果存在,再进一步检查这个方案所需要的辅助量本身,是否也能被廉价地近似出来,而不重新引入本来想要避免的开销。这套两部分的方案(无偏的随机截断,加上为截断方案所需的任意设计权重提供廉价的只用前向传播的代理量),正是本文希望读者能带走、应用到其他 LLM 强化学习效率问题上的可推广经验,而不管 PS-PPO 这个具体的实现最终是否会成为业界标准工具。