Score Centering:训练与采样不一致下的梯度漂移修正

Score Centering:训练与采样不一致下的梯度漂移修正

论文:Score Centering Stabilizes Off-policy Reinforcement Learning
作者:Martin Marek(一作)、Max Ryabinin;Together AI
版本:arXiv:2609.20807v1,2026-09-17
核验日期:2026-09-23
核验方式:对照论文正文、作者仓库的损失实现与实验配置;未复现模型训练。

核心纠正:Score Centering(SC)中心化的是 score function,即对数概率的参数梯度,而不是奖励。 把它实现成 reward -= reward.mean(),会错过论文提出的修正。

一、原报告中需要修正的结论

原报告说法 核验结果与修正
SC 就是 reward 减均值 错误。score 指 \nabla_\theta\log p_\theta;官方代码将奖励组内中心化与 SC 分别实现。
IS 无法消除偏差 错误。满足支持集条件的精确 IS 可以无偏换分布;裁剪、遮罩会引入偏差。
SC 消除全部偏差 过强。去掉 drift 后,仍在采样分布下计算更新。
三种量化设置都达到 50–60%+,其他方法全归零 错误。应逐项读取实验结果,见下表。
一行代码,无额外数据要求 不准确。SC 需要采样端的全词表或 top-k 概率信息。
兼容包括 DPO 在内的所有算法 未验证。仓库支持 DPO,不代表 DPO 已接入 SC;其 DPO 损失没有这项修正。
量化必然崩溃,SC 是唯一解法 无法由这些实验推出。

定义与 IS 的核对见论文第 2.2、4 节;实现对照见 compute_advantages、loss_fn_rl 和 loss_fn_dpo。

二、问题:同一套权重不保证同一个策略

现代 LLM RL 往往让采样与训练使用不同的执行路径。算子、数值精度、量化和 KV cache 处理的差异,可能让同一上下文下的 token 概率不同;异步训练还会引入旧策略生成的数据。

记训练策略为 p_\theta,采样策略为 q。采样来自 q,梯度却通过 p_\theta 计算,这就是需要分析的错位。这里不能直接把未经修正的 off-policy 更新写成真实目标的 \nabla J。

下文给出可独立检查的代数推导。固定一个前缀 h,省略条件 h;所有关于 q 的期望都在这个前缀下计算。R 表示用于加权的回报或已计算好的 advantage,求梯度时视为常量。

三、score、漂移与中心化

定义某个候选 token v 的 score:

s_v=\nabla_\theta\log p_\theta(v\mid h), \qquad \bar{s}_q=\sum_v q(v\mid h)s_v.

它是一个参数梯度向量,不是奖励,也不是 log-prob 数值本身。

由概率归一化可直接得到:

\mathbb E_p[s_v] =\sum_v p_v\nabla_\theta\log p_v =\nabla_\theta\sum_v p_v=0.

但 \mathbb E_q[s_v] 通常不为零。利用协方差定义,未经修正的期望更新为:

G_q=\mathbb E_q[R s_v] =\underbrace{\mathbb E_q[R]\bar{s}_q}_{\text{drift}} +\underbrace{\operatorname{Cov}_q(R,s_v)}_{\text{signal}}.

SC 使用:

\widetilde{s}_v=s_v-\bar{s}_q, \qquad \mathbb E_q[R\widetilde{s}_v] =\operatorname{Cov}_q(R,s_v).

完整分布下,条件期望中的 drift 因此精确抵消。这不表示有限样本没有噪声,也不表示协方差的分布已经从 q 变成 p。该定义对应论文第 4 节。

四、两个手算例子:为什么不能换成 reward 减均值

以下数字是为理解公式构造的例子,不是论文实验。

1. 整组奖励均值为零,局部漂移仍可存在

考虑等概率出现的两个前缀。下一 token 是二元动作 a\in\{0,1\},训练端采用共享参数的 Bernoulli logit,因此 s_a=a-p(1\mid h)。两个前缀使用不同的固定 logit 偏置。

假设奖励只由前缀决定,与当前动作无关;对整个 prompt 的奖励做中心化后,得到下表:

前缀 出现概率 p(1\mid h) q(1\mid h) 条件 advantage \bar{s}_q=q(1\mid h)-p(1\mid h)
h_1 1/2 0.8 0.5 +0.5 -0.3
h_2 1/2 0.2 0.5 -0.5 +0.3

全局 advantage 的期望已经是零,但当前 token 的梯度贡献为:

G_q=\tfrac12(0.5)(-0.3)+\tfrac12(-0.5)(0.3)=-0.15.

当前动作根本不影响奖励,却仍受到更新。SC 在每个前缀减去对应的 \bar{s}_q 后,这部分期望更新变为零。

问题在于,\mathbb E[A]=0 不推出 \mathbb E[A\mid h]=0。如果能获得每个前缀的精确条件奖励均值并将其减去,也可以抵消这一项;普通 batch/group 均值没有提供这个条件信息。

2. 去掉 drift,仍不等于恢复 on-policy 梯度

另取一个单步任务:p(1)=0.8、q(1)=0.5、R=a。对 Bernoulli logit 求导:

G_{\mathrm{SC}}=\operatorname{Cov}_q(a,a-p)=q(1-q)=0.25, \qquad \nabla_\theta\mathbb E_p[R]=p(1-p)=0.16.

两者并不相等。SC 的消漂移性质与“完全无偏地优化训练策略目标”是不同的命题。

相应地,当 q 覆盖 p 的支持集、比值精确可得时,重要性采样满足换测度恒等式:

\mathbb E_q\!\left[\frac{p(y)}{q(y)}R(y)\nabla_\theta\log p(y)\right] =\mathbb E_p[R(y)\nabla_\theta\log p(y)].

这里 y 是完整轨迹。只修正当前 token 的条件概率比,不能自动修正整条轨迹的前缀分布。极端权重带来的方差,以及裁剪后的偏差,需要分别讨论。

五、实现:修改损失,而非修改奖励均值

对于完整词表,下面的教学用损失能产生中心化后的梯度:

L_{\mathrm{SC}} =-\operatorname{sg}(R)\left[ \log p_a-\sum_v\operatorname{sg}(q_v)\log p_v \right].

其中 \operatorname{sg} 表示停止梯度。求导即可检查 -\nabla L_{\mathrm{SC}}=R(s_a-\bar{s}_q)。

# 教学示意:每个位置都有完整的 sampler_probs,最后一维为词表。
# train_logp 来自训练模型的 log_softmax;action_ids 为实际采样 token。
selected_logp = train_logp.gather(-1, action_ids[..., None]).squeeze(-1)
expected_logp = (sampler_probs.detach() * train_logp).sum(dim=-1)
token_loss = -advantage.detach() * (selected_logp - expected_logp)

完整词表的传输和存储可能很贵。官方实现记录采样端 top-k token 的 ID、原始 log-prob,并用训练端分布近似剩余概率质量;不能把 top-k 内概率重新归一化成总和为 1。训练损失还需正确处理生成 token 的 mask。与 IS 组合时,实际采样 token 即使不在 top-k 中,也要保留其采样概率。见作者 README 的 JAX/PyTorch 实现。

源码中,奖励组内中心化位于 compute_advantages;SC 则在 loss_fn_rl 中独立计算 center 并从 score 扣除。启用 SC 却没有 full/top-k 采样分布时,代码会报错。这直接反证了“只改一行 reward 就能部署”的说法。见官方损失源码。

工程上的 top-k 是近似版本,不能把完整分布下的精确消漂移结论直接当作近似实现的恒等式保证。

六、实验结果:按量化设置分别看

30B 实验使用 Qwen3-30B-A3B-Base、INTELLECT-2 math,运行 800 步、最大序列长度 1024、SGD 学习率 0.01;SC 使用 top-128。三档采样端精度和各方法开关可以在官方 30b_math.yaml 中交叉核对。W/A 表示权重与激活精度。

采样端配置 论文报告的训练准确率及行为
FP8 W/A + FP8 KV 未修正 PG 也稳定,达到约 58%。
FP8 W/A + FP4 KV SC 约 52%,TIS 约 51%;PG 在 200 步内崩溃,MIS 后期崩溃。
INT8 W/A + INT4 KV SC 约 30%,TIS 约 12%,其余方法最终低于 5%。

数值来自论文 §5.4。比较使用共享的 REINFORCE+组内中心化目标,仅替换修正规则,并非完整复现各算法配方;论文还明确把人为增强的 TIM 作为压力测试。见§5.1。

严重策略陈旧时,SC 与 TIS/MIS 组合优于单独 SC,见§5.3;配置中的 sampler.staleness=64 和组合开关也可在 06b_quant_stale.yaml 查到。

七、阅读边界

  • 训练准确率不等于泛化收益。 上表不能直接转换成线上能力或端到端吞吐提升。
  • 反馈累积不等于严格线性定律。 原报告的“每步线性累积”和“两周后线上失焦”缺乏对应证据,不作为本笔记结论。
  • 算法支持与框架功能要分开。 作者的 postax 框架也包含 DPO,但其 loss_fn_dpo 没有 SC;是否适用需要针对目标重新推导,不能由目录里存在 train_dpo.py 推出。见框架说明与损失实现。
  • 实际接入要检查采样记录。 已有 reward、advantage 和所选 token 的 log-prob,不代表已经有计算 SC 所需的候选 token 分布。

这项工作的学习价值在于明确了一个可操作的修正对象:检查采样分布下的 score 均值,并修正由它引入的更新漂移。理解它时,应同时保留“逐前缀”“采样分布”和“精确版/近似版”这三个限定。

评论(0)

评论加载中…