Source-linked AI summary
Score Centering Stabilizes Off-policy Reinforcement Learning
Martin Marek, Max Ryabinin
TL;DR
LLM 的 RL 对 sampler–trainer 之间的细微差异(即 TIM)十分敏感,而消除这些差异会降低 rollout 效率。本文指出,drift 是不稳定性的主要机制,并通过 additive score centering 将其抵消。在参数量从 0.6B 到 30B 的量化模型上,score centering 达到或超过 importance-sampling 方法的表现;在严重 staleness 下,两者组合效果最佳。
问题
当 sampler 与 trainer 的输出存在差异时,LLM 的 RL 训练可能变得不稳定,因为在不牺牲 rollout 效率的情况下很难消除 TIM。
方法
本文推导出 score centering:通过减去 sampler 的期望 score,在 TIM 下抵消 drift 的加性修正。
结果
在参数量从 0.6B 到 30B 的 Qwen3 模型上,score centering 在量化条件下达到或超过 importance-sampling 方法的表现,与后者组合时则在严重 staleness 下表现最佳。
结论与局限
在严重量化条件下,drift 足以解释不稳定性;additive score centering 为 importance sampling 提供了实用的互补方案。
结论与局限
Score centering 仍存在基于 sampler 的协方差不匹配;主要结果在短序列上采用了刻意加重的 mismatch,作为较长训练过程中较温和 mismatch 的代理。
Abstract
from arXiv · showhide
Reinforcement learning (RL) of large language models is notoriously sensitive to small differences between training and inference engines, often referred to as the training-inference mismatch (TIM). However, completely eliminating TIM is impractical, as it would come at a major cost to rollout efficiency. In this paper, we show that the instability of RL under TIM is primarily caused by drift: a persistent bias between training and inference engines that accumulates with every training step. We derive an additive "score centering" correction term that stabilizes RL under TIM by canceling drift. When training models from 0.6B to 30B parameters, score centering alone matches or outperforms methods based on importance sampling under quantization, with the gap growing as the mismatch becomes more severe. Because the correction is additive, score centering also composes with importance sampling -- their composition outperforms pure importance-sampling baselines in our staleness experiments.
1 引言
LLM 的 RL 采用基于奖励 rollout 的 policy-gradient 更新,但 sampler 与 trainer 分离的前向传递会产生 TIM,从而导致训练性能下降或奖励崩溃。本文分析了这种不稳定性,并提出 score centering;在报告的实验中,该方法达到或超过了 importance-sampling 稳定化方法。
- RL 的 policy-gradient 训练采样 rollout、分配奖励,并根据奖励加权的 score 计算梯度。
- sampler 与 trainer 分离的前向传递会产生 training-inference mismatch,从而导致训练性能下降或奖励完全崩溃。
- 本文识别出一个 drift 项:当训练策略与采样策略匹配时,该项恰好消失;随后提出 score centering,以稳定 TIM 下的 RL。
- 在 Qwen3-0.6B Countdown 和 Qwen3-30B-A3B-Base INTELLECT-2 实验中,score centering 达到或超过了 importance-sampling 稳定化方法,并且可以与这些方法结合以获得进一步收益。
2 背景与相关工作
TIM 的产生源于:rollout 由 sampler 生成,而 sampler 可能因数值、实现或陈旧性效应而不同于 trainer。Importance sampling 原则上可以精确校正这种不匹配,但实际的 ratio 控制需要在方差与偏差之间权衡,因此 score centering 成为一种从根本上采用加性修正的替代方案。
- 当 rollout 来自 sampler qθ,而梯度通过不同的 trainer pθ 计算时,就会出现 TIM。
- 数值差异、彼此独立的实现、浮点运算的非结合性,以及异步的陈旧 checkpoint,都会产生严重程度不同的 TIM。
- 实际目标是在 TIM 下实现稳定的 RL,同时保留解耦训练与推理所带来的高硬件利用率。
- Importance sampling 通过使用 r = pθ(y)/qθ(y) 对 score 加权来校正 TIM,但原始 ratio 可能变得任意大,从而放大梯度估计的方差。
- 对 importance ratio 设定边界可以稳定实践过程,但必然引入偏差;score centering 不使用 ratio、masking 或 clipping,而是施加加性校正。
3 为什么 RL 对失配敏感
本文区分了离线训练中的奖励不稳定性与 TIM 下在线训练的不稳定性:离线训练偏好非负奖励,而在线训练在使用非负奖励时最不稳定。这表明,在线训练中不断变化的采样器是对失配敏感的关键来源。
- Policy gradient 在二元 +1/0 优势下退化为离线 SFT,因此奖励模式和陈旧性是 SFT 与 RL 之间考察的两个关键差异。
- 受控实验在 Countdown 训练期间扰动 Qwen3-1.7B 的采样器权重,以制造有意的策略失配。
- Figure 2 显示,离线训练只有在使用 +1/0 奖励时才稳定,而 TIM 下的在线训练在使用 +1/0 奖励时最不稳定。
- 离线负奖励不稳定性被归因于无界的负 logprobs 和分布锐化,但本文不再对此展开研究。
- 在线 +1/0 训练类似于从有偏采样器进行蒸馏,该采样器在每一步之后都会刷新权重;这不同于从固定教师模型进行的离线蒸馏。
4 分数中心化
在训练-推理不匹配下,期望策略梯度更新会向采样器方向发生漂移;分数中心化减去采样器的期望分数,以抵消该漂移,包括 off-policy 情况。
- 4 分数中心化: 训练-推理不匹配导致漂移,使训练器逐渐趋向采样器,并在同步过程中不断累积,令全正奖励的在线训练尤其不稳定。漂移取决于平均奖励,而不是哪些 rollout 成功;协方差项则承载学习信号。
- 4 分数中心化: 分数中心化从每个 token 分数中减去采样器的期望分数,使中心化分数在每个前缀处的均值都为零,即使存在不匹配也是如此。由于两种期望都在采样器分布下计算,这会精确抵消漂移。
- 4 分数中心化: 采用分数中心化后,期望更新变为奖励与中心化分数之间的协方差;它与 on-policy 更新的差异仅在于采样分布。协方差恒等式将漂移项与依赖奖励的信号项分离开来。
- 4 分数中心化: Importance sampling 同样可以抵消漂移,但可能增大方差;当较大的比率被裁剪或屏蔽时,还会重新引入漂移,而分数中心化具有确定性且是加性的。由于两种修正彼此独立,分数中心化可以与 importance sampling 组合。
- 4 分数中心化: 分数中心化不同于经典奖励基线,因为它是修正偏差、改变更新均值的方法,而不只是降低方差的工具。在 LLM RL 中,其期望分数可通过对下一个 token 的 logprob 求和精确计算,无需 critic。
- 4 分数中心化: 实际实现只记录 top-k 采样器 logprob,并用训练器分布建模剩余尾部;在测试设置中,k = 128 和 k = 32 的效果与完整中心化相当。该方法以支持 stop-gradient 的标量损失实现,并具有广义 importance-weighted 版本。
5 实验
实验在受控且贴近实际的训练-推理失配条件下比较各种校正方法,结果表明 score centering 始终表现强劲,并且与 importance sampling 结合后在 staleness 条件下尤其有效。
- 5.1 设置: 实验通过 sampler 权重噪声、量化和 staleness 有意放大失配,因为轻微失配需要更长的运行时间才能积累 drift 并区分不同方法。该设置比较了 Countdown 上的 Qwen3-0.6B-Instruct 和 INTELLECT-2 math 上的 Qwen3-30B-A3B-Base,使用共享的训练组件和记录的 sampler 概率。
- 5.2 合成权重噪声: 在最大的合成权重噪声下,只有单独使用 score centering 或将其与 TIS/MIS 结合时才能稳定训练;随着失配加剧,其他方法会更早崩溃。在三种噪声尺度下,失配越大,崩溃越早;DPPO 分别在 steps 160、80 和 20 左右崩溃,而 TIS 在两种较大噪声尺度下分别在 steps 180 和 40 左右崩溃。
- 5.3 量化与 staleness: 在量化条件下,单独使用 score centering 或将其与 TIS/MIS 结合时表现最佳;在严重 staleness 下,各种组合方法占据优势。该比较采用共享的 REINFORCE 目标和组中心化奖励,并分别单独使用或组合使用各校正方法。
- 5.3 量化与 staleness: 在 staleness 条件下,将 score centering 与 importance sampling 结合,被认为可以在 TIS 部分校正采样分布后进一步校正残余 drift。论文建议:当 trainer 在同步步骤之间可能远离 sampler 时,应采用组合方法。
- 5.4 扩展至 30B: 在 30B 规模下,score centering 在严重的 sampler 量化下仍保持稳定:使用 FP4 KV cache 时达到 52%,使用 INT4 KV cache 时达到 30%。在 FP4 KV cache 下,MIS 在后期崩溃,而 score centering 和 TIS 保持稳定;在 INT4 KV cache 下,TIS 达到 12%,其他所有方法最终均低于 5%。
6 结论
论文将 TIM 不稳定性归因于向周期性刷新的采样器漂移,并引入加性 score centering 予以抵消。该方法支持 top-k 近似以及与 importance sampling 的组合,但严重陈旧性和代理不匹配仍是重要边界。
- 6 结论: 在 TIM 下,漂移表现为向采样器进行蒸馏,经周期性同步后不断累积,并且随着不匹配严重程度增加而更早导致崩溃。漂移项取决于平均奖励而非 rollout 成功率,而协方差项承载学习信号。
- 6 结论: Score centering 通过将期望 score 作为加性、无超参数的修正项减去,从而抵消漂移;该修正可表示为标量损失。其广义形式会在与 importance sampling 组合时对加权 score 进行中心化。
- 6 结论: 在轻度 TIM 下,score centering 的表现可匹敌 importance sampling;在严重量化下,它是唯一能够稳定训练的方法,而与 TIS 或 MIS 组合时,在严重陈旧性下表现最佳。对于任意 importance weighting 函数,该修正都能抵消恒定奖励漂移;而 vanilla importance sampling 会使中心化项消失。
- 6 结论: 主要结果在短序列上使用刻意加剧的不匹配,作为轻度不匹配下长期训练的代理;在严重陈旧性下,采样器与训练器之间残留的协方差不匹配也很重要。因此,作者报告称,在该场景中将 score centering 与 importance sampling 组合可获得最佳表现。
- 6 结论: Top-k score centering 根据训练器概率重建采样器尾部,并仅使用头部 token 而非完整词表计算修正。尾部按采样器与训练器尾部质量之比进行重新缩放,同时利用完整词表期望 score 为零这一点恢复期望尾部 score。
- 6 结论: Top-k score centering 的实测开销可忽略不计:在测试的两种模型规模上,k = 128 的运行时间均在基线墙钟时间的 1% 以内。该实现使用 detached sampler probabilities,并支持 TIS 或 MIS 加权函数。
A.4 Top-k 消融实验
在所有测试实验设置中,Top-k score centering 均达到与全词表 score centering 相当的效果,包括失真最严重的量化设置。
- A.4 Top-k 消融实验: k = 32 和 k = 128 在所有测试设置中都达到与完整 score centering 相当的效果,包括 30B 下的 INT8 W/A + INT4 KV。在其他设置中,k = 128 的 top-k 头部平均覆盖超过 99.9% 的 sampler 质量;在最困难的设置中平均覆盖 99.45%,在最差 batch 中覆盖 95.8%。
B.1 训练设置
实验采用基于组中心化奖励的 REINFORCE、SGD、固定学习率 10^-2,以及每个 batch 执行一次优化器更新的批量 rollout。
- B.1 训练设置: 实验采用基于组中心化奖励的 REINFORCE:每个 rollout 的 advantage 都减去其 prompt 组的平均奖励进行中心化。组均值在同一 prompt 的多个 completion 上计算。
- B.1 训练设置: 所有实验中的 SGD 均采用固定学习率 10^-2;该设置训练稳定,性能与 AdamW 相当,同时最多节省 240GB 内存。每个 batch 为每个 prompt 采样八个 completion,并仅执行一次优化器更新。
B.2 修正方法
修正方法对比按照修正作用的位置、比率的裁剪或掩码方式,以及被修正的优势符号进行组织。
- B.2 修正方法: Importance-sampling 方法因 token-level 与 sequence-level 修正、比率裁剪或掩码方式,以及接受修正的优势符号不同而有所差异。DPPO 使用 binary total variation 掩码区域,但仍通过 importance-sampling 权重进行优化。
- B.2 修正方法: Top-k score-centering 消融实验将 k = 32 和 k = 128 与 full-vocabulary score centering 进行对比,随机种子数量标注在图例中。该消融可视化结果见 Figure 5。
- B.2 修正方法: 所有修正都应用于同一目标函数,其中 score centering 使用 k = 128 的 top-k 近似实现。PG 是未经修正的基线,其他方法采用其原论文或 verl 中的默认参数,未进行进一步调参。
B.3 计算资源
实验使用数量不一的随机种子,并报告主要在 8× NVIDIA H100 SXM GPU 节点上复现图表所需的大致计算资源。
- 0.6B 和 1.7B 模型在 Figures 2–4 的每项实验和每种校正方法中均使用 3 个随机种子。
- 30B 模型在部分实验中使用较少的随机种子,因为其训练成本显著更高。
- Table 2 汇总了每条曲线使用的随机种子数量,以及主要在 8× NVIDIA H100 SXM GPU 节点上复现图表所需的大致计算资源。