Score Centering Stabilizes Off-policy Reinforcement Learning
Martin Marek, Max Ryabinin
cs.LG
2026-09-18
采样和训练两套引擎对不齐时,policy gradient会把trainer往sampler蒸馏,误差逐步放大。减去每个前缀的期望score后,INT8+INT4下30B准确率30%,TIS仅12%。
LLM 的强化学习要跑两套前向:推理引擎采样,训练引擎算梯度。两者数值对不齐,就是 training-inference mismatch(TIM)。来源包括量化、内核、浮点归约顺序、异步过期。完全消掉 TIM 会牺牲 rollout 效率。SFT 对离线数据很稳,RL 却会 reward collapse。现成补丁几乎全是 importance sampling,再 clip 或 mask 压方差,clip 又把偏差送回来。
把期望 policy gradient 拆开。奖励乘 score 的期望,等于平均奖励乘平均 score,再加上奖励和 score 的协方差。前一项叫 drift,后一项才是信号。
On-policy 时平均 score 必须是 0:概率归一化,对参数求导得到 0。有 TIM 时 token 从 sampler q 来,score 在 trainer p 上算,平均 score 一般不是 0。drift 不含「哪个 token 更好」的信息,方向是把 trainer 往 sampler 做 SFT,强度跟该前缀的平均奖励成正比。sampler 是 trainer 的有偏副本,权重还会同步回去,偏差正反馈。这解释了为什么在线训练里全正奖励最不稳:每一步都在蒸馏一个跟着自己走的老师。
Group centering 消不掉 drift。某个前缀的期望 advantage 不是 0,正好是有学习信号的位置 drift 最大。
Score centering 从每个 score 里减掉 sampler 下的期望 score。加性修正,给定前缀就确定,跟抽到哪个 token 无关,也没有超参。期望更新变成协方差项,和 on-policy 只差一件事:协方差在 q 上量,不在 p 上量。没有 TIM 时修正项恰好为 0。
全词表期望存不下。Qwen3 词表约 15.2 万,长序列要 TB 级内存。实现只记 sampler 的 top-k logprob,默认 k=128,k=32 在他们测过的设置里也对齐全词表;尾巴用 trainer 分布按尾巴质量缩放,写成带 stop-gradient 的标量 loss。修正是加性的,可以和截断或掩码 importance sampling(TIS/MIS)组合。
共用目标:REINFORCE 加 group-centered reward,只换修正项。Qwen3-0.6B 跑 Countdown,Qwen3-30B-A3B-Base 跑 INTELLECT-2 数学子集。
合成权重噪声越大,崩得越早。最大噪声下只有 score centering,以及它和 TIS/MIS 的组合能训完。
更真实的两档 TIM:只量化 sampler;推理引擎每 64 步才更新一次。量化下单独用 score centering 就最好。过期下要叠 TIS 或 MIS,因为协方差还在 q 上量,IS 先把分布往 p 拉一点。
30B 三档量化:
| 设置 | score centering | TIS | 朴素 PG |
| FP8 sampler | 与朴素 PG都稳 | — | 训练准确率 58% |
| KV 量化到 FP4 | 52% | 51% | 200 step 内崩 |
| INT8 权重 + INT4 KV | 30% | 12% | 其余方法低于 5% |
PPO 和 DAPO 扛得住过期,扛不住量化和权重噪声。它们的 clip 是为策略移动设计的,数值误差不长那样的 importance ratio。
崩的原因被收成一句:会跟着学生走的老师,不是 mismatch 的绝对值。离线蒸馏稳、在线 RL 不稳,对上这个机制。取消 drift、不用 importance ratio,在重度量化下已经够稳住,说明 drift 就是主因。
工程上这是几行 loss,和现有 IS 正交。重度量化 rollout 或异步长 episode,可以先上 score centering;过期严重再叠 TIS/MIS。不必为了对齐数值去换 batch-invariant kernel,或把训练推理都锁同一精度。
修正消掉 drift,剩下的更新仍在 sampler 分布上量协方差,过期严重时必须叠 IS。头条数字来自故意加重的 TIM 和短序列:Countdown 最长 512 token,30B 最长 1024,用来代理「更轻 TIM、更长训练」。MoE 没有 replay 路由。没有在生产级轻度 TIM、数千 step 上验证「drift 慢慢攒起来」这条外推。