发表时间: 2026-09 · arXiv:2609.20807 (Together AI)
原文: https://arxiv.org/abs/2609.20807
作者/机构:Martin Marek & Max Ryabinin, Together AI
一句话结论 本文发现大模型强化学习中训练与推理引擎不匹配导致的不稳定源于梯度漂移,并提出了一种名为 Score Centering 的附加校正项,在严重量化和延迟更新等场景下,该方法能有效消除漂移,其单独使用或与重要性采样结合均能显著稳定训练并超越现有基线。
要解决什么问题 大语言模型的强化学习(如 PPO、GRPO 等基于策略梯度的方法)通常需要两个前向传播:采样器(推理引擎)生成数据,训练器(训练引擎)计算梯度。理想情况下两者是同一个模型,但实际工程中,为了追求极致的生成效率,采样器和训练器往往存在差异。这种差异被称为训练-推理不匹配(TIM)。TIM 的来源包括:两者使用不同的精度(如采样器用 INT8,训练器用 BF16)、不同的底层算子导致浮点数计算顺序不同、以及为了提高硬件利用率而采用的异步生成导致采样器使用的是旧版本权重(即更新延迟)。当采样策略和训练策略不一致时,标准的策略梯度更新会产生一个漂移项。因为在线强化学习中,采样器本质上是训练器的一个带有偏差的副本,并且会定期从训练器同步权重。这种漂移相当于让训练器不断向带有偏差的采样器进行知识蒸馏。随着训练步数的增加,这种偏差会在反馈循环中不断累积放大,最终导致模型输出崩溃。虽然现有的重要性采样方法可以通过对梯度进行加权来纠正这种不匹配,但重要性权重在罕见 token 上会变得极大,导致梯度方差剧增。为了控制方差,工程上通常会对权重进行裁剪或掩码处理,但这又不可避免地重新引入了偏差,无法从根本上解决严重不匹配下的崩溃问题。
怎么做的 为了彻底消除梯度漂移而不引入高方差,作者提出了一种名为 Score Centering(分数居中)的方法。其核心思路是:即使在训练-推理不匹配的情况下,我们也希望在任何前缀下,期望的策略梯度更新为零(就像在同策略训练中一样)。具体而言,标准的策略梯度更新包含一个分数(即对数概率的梯度)。作者通过协方差恒等式证明,当采样器 $q$ 和训练器 $p$ 不一致时,期望的策略梯度更新可以分解为漂移项和信号项。为了消除漂移,Score Centering 直接从每个 token 的分数 $s_{y_t}$ 中减去该分数在采样器分布下的期望值 $\bar{s}$,得到居中后的分数 $\tilde{s}_{y_t}$:
其中 $v$ 遍历整个词表。这样一来,居中后的分数在采样器分布下的期望严格为零,从而在数学上完全抵消了漂移项。与重要性采样使用随机的乘性权重不同,Score Centering 是一个确定性的加性校正项,它不依赖于具体采样到了哪个 token,因此不会增加梯度方差,也就不需要任何裁剪或掩码操作。在工程实现上,计算完整词表的期望分数会导致显存溢出。因此,该方法采用了一种 Top-k 近似策略:在采样时仅记录采样器输出概率最高的 $k$ 个 token 的对数概率(如 $k=128$),而对于长尾部分,则使用训练器的概率分布进行缩放填补,使其总概率和为一。这种近似使得校正项的计算仅需在头部 token 上进行,计算开销几乎可以忽略不计。此外,由于 Score Centering 是加性校正,它与乘性的重要性采样在机制上完全正交,因此可以无缝组合使用,在消除漂移的同时进一步修正采样分布。
效果如何 实验在两个规模上进行:一是使用 0.6B 参数的 Qwen3 模型在 Countdown 任务上训练,二是使用 30B 参数的 Qwen3-30B-A3B-Base 混合专家模型在 INTELLECT-2 数学数据集上训练。为了在有限算力下观察到显著差异,作者故意放大了不匹配的严重程度,设置了三种极端场景:人为添加权重噪声、对采样器进行严重量化(如 INT8 甚至 FP4 的 KV 缓存)、以及极端的更新延迟(采样器每 64 步才更新一次)。对比的基线方法涵盖了当前主流的路线,包括不加校正的朴素策略梯度(PG),以及代表重要性采样路线的截断重要性采样(TIS)、掩码重要性采样(MIS)、PPO 和 DAPO。在 30B 模型且采样器使用 FP8 精度的轻度不匹配下,所有方法都能稳定训练并达到约 58% 的准确率。但当 KV 缓存被量化为 FP4 时,PG 在 200 步内崩溃,MIS 在训练后期崩溃,而 Score Centering 和 TIS 依然保持稳定,分别达到 52% 和 51% 的准确率。在最极端的 INT8 采样器加 INT4 KV 缓存设置下,Score Centering 达到了 30% 的准确率,TIS 仅为 12%,其余所有基线方法均跌破 5%。在 0.6B 模型的极端延迟更新实验中,单独的 Score Centering 表现不如预期,但将其与 TIS 或 MIS 组合使用后,取得了全场最优的稳定性。作者也坦承了该方法的局限性:Score Centering 虽然消除了漂移,但剩余的更新项衡量的是奖励与分数在采样器分布下的协方差,而非训练器分布下的协方差。这种分布偏差在严重的延迟更新场景下会显现出来,这也是为什么此时必须将其与重要性采样组合使用才能达到最佳效果。此外,为了加速实验,核心结论多基于短序列和人为放大的极端不匹配场景得出。
当前大语言模型(LLM)的强化学习(RL)对训练引擎和推理引擎之间的微小差异极为敏感,这种现象通常被称为训练-推理不匹配(Training-Inference Mismatch, TIM)。然而,在实践中完全消除TIM是不切实际的,因为这将对模型的采样生成(Rollout)效率造成巨大代价。本文通过研究发现,在TIM下强化学习的不稳定性主要是由“漂移(Drift)”引起的:这是一种在训练和推理引擎之间持续存在的偏差,并且会随着每一次训练步骤不断累积。为了解决这一问题,本文推导并提出了一种加性“得分中心化(Score Centering)”校正项,通过抵消漂移来稳定TIM下的强化学习。
在对0.6B到30B参数规模的模型进行训练时,仅仅使用得分中心化方法,就能在量化场景下达到或超越基于重要性采样(Importance Sampling)的方法,并且随着不匹配程度的加剧,这种性能差距会进一步扩大。此外,由于该校正项是加性的,得分中心化还可以与重要性采样方法结合使用——在实验的延迟更新(Staleness)场景中,它们的组合表现优于纯粹的重要性采样基线方法。
训练-推理不匹配(TIM)的来源:在标准的策略梯度等式中,仅当用于计算梯度的策略与生成轨迹(Rollouts)的策略相同时才成立。然而在实践中,轨迹是从采样器 $q_\theta$ 生成的,而梯度是通过训练器 $p_\theta$ 计算的。因此,当 $q_\theta \neq p_\theta$ 时,就会发生训练-推理不匹配(TIM)。
严重TIM的成因:TIM有许多不同的来源,其严重程度各异。在最坏的情况下,训练和推理引擎可能运行在不同的精度下,或者基于两个独立的代码库,每个代码库都可能有不同的(难以察觉的)漏洞,从而导致输出存在微妙的差异【5,fix(attention): handle extreme negative logits in masked softmax + 2026 + GitHub + https://github.com/flashinfer-ai/flashinfer/pull/4401】、【16 ,When speed kills stability: Demystifying RL collapse from the training-inference mismatch + 2025 + Blog + https://richardli.xyz/rl-collapse】 。
轻度TIM的成因:在较轻微的情况下,即使两个引擎都是正确的且使用相同的GPU内核,由于浮点运算的非结合律特性,它们仍可能产生不同的输出。因为采样是自回归的,而训练通常在序列维度上并行化,训练和推理引擎可能会以不同的输入形状调用相同的内核,改变了归约(Reductions)的顺序,从而导致微小的浮点数差异。虽然批次不变(Batch-invariant)内核消除了对输入形状的依赖,但这需要大量的工程工作,并会导致更差的GPU利用率【9,Defeating nondeterminism in LLM inference + 2025 + Blog + https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/】。将两个引擎从bf16精度切换到fp16精度也能极大减少不匹配【27 ,Defeating the training-inference mismatch via FP16 + 2025 + arXiv + https://arxiv.org/abs/2510.26788】,但无法完全消除它 。
更新延迟导致的TIM:最后,即使训练器和生成器使用完全相同的批次不变内核,TIM仍可能因更新延迟(Update staleness)而发生。为了在RL期间最大化GPU利用率,需要将训练和推理引擎解耦,使它们异步运行【6,AReaL: A large-scale asynchronous reinforcement learning system for language reasoning + 2025 + arXiv + https://arxiv.org/abs/2505.24298】,并通过连续批处理 (Continuous batching)提供轨迹服务【40,Orca: A distributed serving system for transformer-based generative models + 2022 + OSDI】。这意味着单个训练批次,甚至单个轨迹的不同子序列,可能是从不同的模型检查点生成的【25,PipelineRL: Faster on-policy reinforcement learning for long sequence generation + 2025 + arXiv】。在智能体编码(Agentic coding)等长上下文环境中,延迟变得最为严重,因为某些回合可能在几分钟内完成,而其他回合可能持续数小时甚至数天【11,Single-rollout asynchronous optimization for agentic reinforcement learning + 2026 + arXiv + https://arxiv.org/abs/2607.07508】。从硬件利用率的角度来看,通过算法来稳定带有过期轨迹的训练,比试图完全消除延迟更为可取【6 ,AReaL: A large-scale asynchronous reinforcement learning system for language reasoning + 2025 + arXiv + https://arxiv.org/abs/2505.24298】 。
设计目标:本文的目标是提出一种在TIM下稳定RL训练的方法,以在确保训练稳定的同时实现高硬件利用率。
重要性采样校正原理:假设在策略梯度等式中,轨迹来自于采样器 $q_\theta$ 而不是训练器 $p_\theta$,且 $q_\theta \neq p_\theta$。那么我们可以使用重要性采样(IS)精确地校正训练-推理不匹配:
重要性采样的方差与偏差权衡:上述等式是为完整序列编写的;在实践中,该比率通常在每个词元(Token)级别应用【42,Group sequence policy optimization + 2025 + arXiv + https://arxiv.org/abs/2507.18071】。无论采用哪种方式,重要性采样在校正TIM的同时都付出了方差增加的代价。加权得分被乘以重要性比 率 $r = \frac{p_\theta(y)}{q_\theta(y)}$,该比率在罕见词元上可能会取任意大的值,从而膨胀梯度估计的方差【12,Truncated importance sampling + 2008 + Journal of Computational and Graphical Statistics】、【24,Monte Carlo Theory, Methods and Examples + 2013 + Book】。由于使用原始重要性比率进行训练是不稳定的,实际中使用的方法会对重要性比率进行有界约束,这不可避免地引入了偏差。
现有重要性采样方法的分类:附录表1显示,大多数流行的校正方法都基于重要性采样,并且根据重要性比率被裁剪或掩码的区域,可以归入一个简单的网格分类中。一个显著的例外是DPPO,它使用基于二元总变差(Binary total variation)的掩码区域——但其目标函数仍然使用了IS比率,就像网格中的其他所有方法一样。超越该表格的范畴,Ye等人【39,Adaptive layerwise perturbation: Unifying off-policy corrections for LLM RL + 2026 + arXiv + https://arxiv.org/abs/2603.19470】向训练器的隐藏状态注入可学习的扰动,并使用扰动后的策略作为重要性比率的分子——以此减少比率的重尾分布而不是裁剪它们——但其仍然依赖于重要性采样。相比之下,得分中心化(我们在后续章节介绍的方法)不使用重要性比率,不使用掩码,也不使用裁剪——它是一个加性校正项,其工作原理与此网格中的每种方法都有根本的不同 。
RL与SFT的敏感性对比:为了理解为什么RL对TIM如此敏感,了解RL与监督微调(SFT)的不同之处是有帮助的。正如我们之前指出的,SFT在面临严重得多的不匹配时依然稳定——即在完全由不同模型离线生成的数据上进行训练,且无需任何重要性采样校正【10,Distilling the knowledge in a neural network + 2015 + arXiv + https://arxiv.org/abs/1503.02531】、【2 ,DeepSeek-R1: Incentivizing reasoning capability in LLMs via reinforcement learning + 2025 + arXiv + https://arxiv.org/abs/2501.12948】。换句话说,SFT对延迟或内核的选择完全不敏感 。
策略梯度退化为SFT的条件:至关重要的是,在一个具有二元奖励的任务中,当优势(Advantages)为 $+1/0$(即奖励未被中心化或归一化)且训练是完全离线(即采样器从不更新)时,策略梯度就退化为SFT。因此,我们可以认为奖励模式(Reward mode)和延迟(Staleness)是SFT和RL之间的两个核心区别。我们希望确定这两种机制中哪一个是导致对TIM敏感性增加的原因。
受控TIM实验设置:我们通过一个具有故意策略不匹配的离线/在线训练玩具实验来测试这一点。我们在Countdown数据集【7,Stream of search (SoS): Learning to search in language + 2024 + arXiv + https://arxiv.org/abs/2404.03683】上训练Qwen3-1.7B【37 ,Qwen3 technical report + 2025 + arXiv + https://arxiv.org/abs/2505.09388】,并人为地对采样器的权重施加一个相对于训练器的小扰动,从而创建受控的TIM 。
奖励模式与稳定性的关系分析:图2显示,在这种设置下,离线训练仅在 $+1/0$(非负)奖励下稳定,而在线训练的情况恰恰相反——组内中心化奖励(Group-centered rewards)【33,DeepSeekMath: Pushing the limits of mathematical reasoning in open language models + 2024 + arXiv + https://arxiv.org/abs/2402.03300】由于在一个组内混合了正负值,表现最为稳定,而对于离线训练最稳定 的 $+1/0$ 奖励,在在线训练中却最不稳定。因此,我们假设存在两种独立的机制在起作用:一种机制导致具有负奖励的离线训练变得不稳定,另一种机制导致具有正奖励的在线训练变得不稳定。
离线负奖励不稳定的机制说明:本文的主要重点是在线强化学习,因此我们不再深入探讨负奖励下离线训练的不稳定性。Ren & Sutherland【31,Learning dynamics of LLM finetuning + 2025 + ICLR】认为,这种不稳定性背后的主要机制是负奖励的无界性结合分布锐化(Distribution sharpening)。对数概率在下方是无界的,因此在带有负奖励的离线训练期间,它们可能会发散到 $-\infty$。这种效应在仅采样高概率词元的在线训练中是不存在的。这种不对称性有充分的文献记录:仅对正样本进行训练(即蒸馏或拒绝微调)即使在离线状态下也是完全稳定的【10,Distilling the knowledge in a neural network + 2015 + arXiv + https://arxiv.org/abs/1503.02531】、【2 ,DeepSeek-R1: Incentivizing reasoning capability in LLMs via reinforcement learning + 2025 + arXiv + https://arxiv.org/abs/2501.12948】,而负梯度可能会拉低正确响应的概率【3 ,On the effect of negative gradient in group relative deep reinforcement optimization + 2025 + NeurIPS + https://arxiv.org/abs/2505.18830】——尽管在同策略(On-policy)下,它们携带了有用的学习信号【43 ,The surprising effectiveness of negative reinforcement in LLM reasoning + 2025 + arXiv + https://arxiv.org/abs/2506.01347】 。
在线正奖励不稳定性的本质:在在线训练中,我们观察到与离线训练相反的行为——仅使用正奖励是最不稳定的设置。因此,我们假设这种不稳定性必然源自训练过程的在线特性。需要注意的是,在TIM下使用 $+1/0$ 奖励的在线训练本身就是一种蒸馏形式:训练器拟合采样器生成的成功轨迹,而采样器是训练器的一个带有偏差的副本,其权重在每一步后都会从训练器处刷新。与离线蒸馏的唯一区别在于,教师模型(采样器)会随着学生模型(训练器)移动,因此这必定是不稳定性的来源。
参数与符号定义:来自轨迹 $y$ 的更新汇总了其各个词元的得分,每个得分由轨迹的奖励(或优势)$R$ 加权。我们检查一个词元 $y_t$ 在给定其前缀 $y_{<t}$ 的条件下的预期贡献。采样器 $q$ 生成词元 $y_t$;$v$ 索引词汇表。$p_v$ 和 $q_v$ 分别是 $v$ 在训练器 $p$ 和采样器 $q$ 下的下一个词元概率,$s_v = \nabla_\theta \log p_v$ 是其得分,而 $\bar{s} = \sum_v q_v s_v$ 是在采样器下的预期得分。$\mathbb{E}_q$ 和 $\mathbb{E}_p$ 是对前缀之后轨迹其余部分的期望。</p>
同策略下的恒定奖励分析:考虑在一个具有恒定 $+1$ 奖励的环境中使用朴素策略梯度进行RL训练。理论上,不应该发生任何学习——奖励是恒定的,所以没有来自环境的“信号”。实际上,在同策略下,每个前缀的预期梯度确实为零:$\mathbb{E}_p[R s_{y_t}] = R \mathbb{E}_p[s_{y_t}] = R \cdot 0$,因为对于任何分布都有 $\mathbb{E}_p[s_{y_t}] = \sum_v p_v \nabla_\theta \log p_v = \nabla_\theta \sum_v p_v = \nabla_\theta 1 = 0$。
TIM下的梯度漂移推导:然而,在TIM下,预期梯度非零:词元是从采样器 $q$ 采样的,而得分是通过训练器 $p$ 计算的,因此通常情况下 $\bar{s} = \mathbb{E}_q[s_{y_t}] \neq 0$。利用协方差恒等式,我们可以将前缀处的预期策略梯度更新分解为“漂移(Drift)”项和“信号(Signal)”项:
``
漂移项的蒸馏效应解析:漂移是TIM的产物,它起到了向采样器蒸馏的作用。漂移项不包含关于哪些轨迹成功的信息:它仅通过奖励的均值 $\mathbb{E}_q[R]$ 来依赖奖励。只有协方差项才能看出哪个词元导致了哪个奖励。漂移项的实际作用由其方向 $\bar{s}$ 决定:预期得分是以采样器为教师的交叉熵(SFT)损失的负梯度,因此朴素策略梯度在每个前缀处将训练器向采样器进行蒸馏,其缩放比例为该前缀处的预期奖励。这纯粹是TIM的产物:在同策略下,训练器已经与采样器匹配,因此 $\bar{s} = 0$,漂移项消失。向固定的教师进行蒸馏是无害的——它会收敛到该教师的最大似然拟合。然而,采样器并不是固定的:它是训练器的一个带有偏差的副本(例如被量化或延迟),因此每一步都会将训练器推向采样器;然后训练器的权重被同步回采样器,误差在一个反馈循环中不断复合而不是收敛。这解释了为什么在图2中,带有全正奖励的在线训练是最不稳定的设置。Dong等人【4,Probing RLVR training instability through the lens of objective-level hacking + 2026 + ICML + https://openreview.net/forum?id=KlGj06E8Wa】同样将系统性偏差确定为TIM下不稳定性的根本原因,并表明它通过正反馈循环不断复合 。
组内中心化无法消除漂移:组内中心化(Group centering)【33,DeepSeekMath: Pushing the limits of mathematical reasoning in open language models + 2024 + arXiv + https://arxiv.org/abs/2402.03300】使得一个提示词的各个轨迹上的优势总和为零,但漂移出现在单个前缀处,而在给定前缀处的预期优势并不为零:一个极有可能导致正确答案的前缀具有正的预期优势,而一个已经包含错误的前缀具有负的预期优势。因此,漂移恰好在携带学习信号的前缀处非零——无论所讨论的词元是否影响奖励,训练器在具有正预期优势的前缀后被拉向采样器,在具有负预期优势的前缀后被推离采样器。 与 $+1/0$ 奖励相比,组内中心化确实缩小了漂移,这与它在图2中是最稳定的在线设置相一致,但它并没有消除漂移。
得分中心化消除漂移的原理:上述观察直接激发了我们的方法,称为得分中心化(Score Centering):即使在TIM下,我们也希望每个前缀处的预期得分为零,就像在同策略下一样。我们只需从每个得分中减去采样器下的预期得分即可实现这一点,用中心化得分 $\tilde{s}_{y_t}$ 替换 $s_{y_t}$:
得分中心化下的期望更新分析:采用得分中心化后,预期的策略梯度更新变为:
如公式(4)所定义的中心化得分在采样器下具有零均值(即$\mathbb{E}_q[\tilde{s}_{y_t}] = \bar{s} - \bar{s} = 0$),因为两个期望都是在相同的分布下获取的——漂移在每个前缀处被精确抵消,即使在异策略(Off-policy)下也是如此。公式(5)表明,采用得分中心化的预期更新等于预期的同策略更新,唯一的区别在于测量协方差所在的分布——在同策略和具有重要性采样的TIM下,协方差是在训练分布 $p_\theta$ 下测量的,而得分中心化是在采样分布 $q_\theta$ 下测量它。在没有TIM的情况下,即当 $p_\theta = q_\theta$ 时,得分中心化是一个空操作(No-op),因为公式(4)中的校正项恰好为零。
得分中心化与重要性采样的对比:重要性采样与得分中心化一样,也能精确抵消漂移。然而,重要性采样使用一个随机的乘法校正项,这意味着它增加了梯度方差,并且其值取决于采样的词元。对于罕见词元,重要性比率可能非常大,这就是为什么几乎所有实际实现都以某种方式裁剪或掩码大重要性比率(表1),从而重新引入了漂移。相比之下,得分中心化是一个加性校正项,给定前缀它是确定性的——它不依赖于采样的词元,并且可以精确评估。由于得分中心化和重要性采样以独立的方式校正TIM,这两种方法可以相互组合。
与经典基线方法的联系与区别:从策略梯度中减去一个零均值量是一个经典的方差缩减思想:奖励基线(Reward baselines)【35,Simple statistical gradient-following algorithms for connectionist reinforcement learning + 1992 + Machine Learning】和得分函数控制变量(Score-function control variates)【30,Black box variational inference + 2014 + AISTATS】都依赖于恒等式 $\mathbb{E}_p[s_{y_t}] = 0$。得分中心化是一种偏差校正而不是方差缩减:减去的向量给定前缀是确定性的,其选择是为了移动更新的均值,而在同策略下,奖励基线改变的是方差而不是均值。在异策略下,等于前缀处预期奖励的奖励基线(一个精确的价值函数)在期望上也能抵消漂移;得分中心化在没有评论家(Critic)模型的情况下获得了相同的预期更新。据我们所知,之前的任何方法都没有减去预期得分本身:在同策略下它为零,而在经典的异策略RL中,它需要对动作空间进行难以处理的期望计算,因此异策略方法转而依赖重要性采样。在TIM下的LLM RL训练的独特之处在于,预期得分不仅非零,而且可以精确计算,即作为下一个词元对数概率的总和。
实现机制与Top-k近似:存储采样器的完整下一个词元分布在计算上过于昂贵。因此,我们仅记录其前 $k$ 个对数概率(例如 $k=128$),并用训练器的分布来建模尾部,进行重新缩放以匹配采样器的尾部概率质量。我们将得分中心化实现为一个标量损失函数:
其中 $\rho$ 是采样器与训练器尾部质量的比率,$\mathrm{sg}$ 表示停止梯度(Stop-gradient)。在测试的每一个设置中,使用 $k=128$ 甚至 $k=32$ 都能达到与完整得分中心化相匹配的效果。以下代码片段说明了单个词元的最小实现:
import jax.numpy as jnp
from jax.lax import stop_gradient
def score_centering_loss(train_logp, samp_logp, topk_ids, sampled_token, advantage):
train_head_logp = train_logp[topk_ids]
tail_mass_ratio = (1 - jnp.exp(samp_logp).sum()) / (1 - jnp.exp(train_head_logp).sum())
head_prob_residual = jnp.exp(samp_logp) - tail_mass_ratio * jnp.exp(train_head_logp)
logp_correction = (stop_gradient(head_prob_residual) * train_head_logp).sum()
return -advantage * (train_logp[sampled_token] - logp_correction)
代码功能解析:首先通过索引获取训练器在Top-k词汇上的对数概率 train_head_logp。然后计算尾部质量比率 tail_mass_ratio,即采样器尾部概率之和除以训练器尾部概率之和。接着计算头部概率残差 head_prob_residual,即采样器头部概率减去按比率缩放后的训练器头部概率。通过将停止梯度的残差与训练器头部对数概率相乘并求和,得到对数概率的校正项 logp_correction。最后,损失函数返回优势值乘以(采样词元的训练对数概率减去校正项)的相反数。
数据集与模型配置:实验主要在两个数据集和模型上进行验证。其一是在 Countdown 数据集上训练 Qwen3-0.6B-Instruct 模型;其二是在 INTELLECT-2 数据的数学子集上训练 Qwen3-30B-A3B-Base 模型。
硬件与软件配置:实验主要在配备 $8 \times$ NVIDIA H100 SXM GPU 的节点上运行。代码实现使用了自定义的 JAX 采样器。优化器选用 SGD,固定学习率为 $10^{-2}$。
训练超参数配置:所有实验均使用带有组内中心化奖励的 REINFORCE 算法目标函数。在 0.6B 模型实验中,每个批次包含 64 个提示词(共 512 个序列),每个提示词采样 8 个完成结果,最大序列长度为 512。在 30B 模型实验中,每个批次包含 16 个提示词(共 128 个序列),最大序列长度为 1024。每个批次执行一次 SGD 优化器更新。
不匹配(TIM)设置:为了在可控且有限的计算资源内观察到显著差异,实验设计了三种故意放大的严重 TIM 场景:
合成权重噪声的实验结果:图 3 展示了在三种不同噪声尺度下的实验结果。得分中心化(SC)、截断重要性采样(TIS)和掩码重要性采样(MIS)表现最佳,它们的组合($\mathrm{MIS} + \mathrm{SC}$ 和 $\mathrm{TIS} + \mathrm{SC}$)同样出色。随着噪声增加,训练更早崩溃(例如 DPPO 分别在第 160、80 和 20 步崩溃)。在最严重的噪声下,只有得分中心化(单独使用或与 TIS/MIS 组合)能够保持稳定训练。这证实了漂移会随着训练不断累积,更大的不匹配会导致更早的崩溃。
量化与延迟的实验结果:图 4 展示了在更现实的 TIM 来源(量化和延迟)下的表现。在量化设置(INT8 采样器)下,结果与权重噪声实验类似,得分中心化(单独或组合)表现最佳。然而,在延迟设置(每 64 步更新一次)下,得分中心化与 TIS/MIS 的组合表现最好,优于基础的得分中心化。这是因为在这种情况下,SC 测量的协方差是在采样器分布下而非训练器分布下,TIS 部分校正了采样分布,而 SC 移除了剩余的漂移。值得注意的是,PPO 和 DAPO 在延迟下存活但在量化和权重噪声下崩溃,这与它们设计的裁剪机制一致(针对策略移动产生的比率,而非数值误差)。
扩展到 30B 模型的实验结果:图 1 验证了结果在大规模模型上的有效性。在 FP8 采样器下,甚至未校正的策略梯度(PG)也能稳定训练,达到 58% 的训练准确率。当 KV 缓存量化为 FP4 时,PG 在 200 步内崩溃,MIS 在训练后期崩溃,而得分中心化(52%)和 TIS(51%)保持稳定。在极端量化(INT8 采样器和 INT4 KV 缓存)下,得分中心化达到 30% 准确率,TIS 为 12%,其他所有方法均低于 5%。
本文解释了为什么 LLM 的 RL 训练对训练-推理不匹配(TIM)如此敏感。在 TIM 下,策略梯度更新包含一个漂移项,该项起到向采样器蒸馏的作用。由于采样器是训练器的有偏副本并定期同步,这种偏差在反馈循环中不断复合。本文提出,仅通过消除漂移而无需使用重要性比率,就足以在严重量化下稳定训练,确立了漂移是不稳定性的根本原因。
本文提出的“得分中心化”是一种实用的方法,它通过减去预期得分来抵消漂移。它是一种无超参数的加性校正项,可以表示为标量损失,并且能够与 TIS 和 MIS 等重要性采样方法组合使用。为了解决精确计算的开销问题,本文引入了 Top-$k$ 近似方法,在实验中完美匹配了精确计算的性能。在轻度 TIM 下,得分中心化与重要性采样方法表现相当;在严重量化下,它是唯一能稳定训练的方法;在严重延迟下,它与 TIS 或 MIS 的组合表现最佳。未来的工作需要注意,尽管得分中心化消除了漂移,但剩余更新测量的协方差是在采样器而非训练器下进行的,这在严重延迟下具有一定影响。
Top-k 近似的计算成本与实现:正如目前所描述的,完整的得分中心化对于大多数实际应用来说过于昂贵。根据公式(4)描述的计算得分中心化校正项(在采样器的逐词元分布上对得分求期望)需要存储每个采样词元的完整输出对数概率。例如,Qwen3模型的词汇表大小为152K,因此在fp32精度下仅存储单个词元的完整对数概率就需要约608KB的内存。假设一个包含1024个轨迹的批次,每个轨迹序列长度为32K,存储完整的对数概率将需要令人望而却步的20TB内存。因此,在我们的整个实验中,我们仅存储采样器分布的Top-$k$近似值,其中 $k=128$。然而,我们并非仅仅在Top-$k$分布上计算预期得分,而是仍然尝试近似采样器的完整分布,使用来自训练器的对数概率来填充尾部。我们使用 $k=128$ 作为保守的默认值;图5显示,在所有测试设置中,$k=128$ 甚至 $k=32$ 的表现都与完整的得分中心化相当。
完整分布的重构逻辑:接下来的问题是如何仅从采样器的Top-$k$对数概率重构其完整分布。令 $\hat{q}$ 表示我们对采样分布的近似。我们从采样器获取Top-$k$对数概率,并使用从训练器获取的对数概率对尾部进行建模,重新缩放以使整个分布的总和为一。将Top-$k$词元集合表示为头部 $H$,其余词元表示为尾部 $T$:
标量 $\rho$ 是采样器尾部质量与训练器尾部质量的比率——它重新缩放训练器的尾部,使得 $\hat{q}$ 的总和为一。
预期得分的降本计算:为了计算在这个重构的 $\hat{q}$ 分布下的预期得分,不需要对整个词汇表求和。相反,我们精确地计算从 $q$ 获取的头部的期望,并将从 $p$ 获取的尾部的期望计算为零减去头部的期望,这利用了在完整词汇表上的预期得分总和为零的事实:
因此,中心化项仅涉及 $k$ 个头部词元,尽管建模的尾部覆盖了整个词汇表。
Top-k 方法的性能开销:Top-$k$ 得分中心化的成本可以忽略不计:在匹配的硬件上,使用 $k=128$ 的运行对于0.6B和30B模型都在基线方法的实际挂钟时间(Wall-clock time)的 $1\%$ 以内完成。我们的自定义JAX采样器在解码过程中计算Top-$k$对数概率;vLLM和SGLang原生暴露了Top-$k$对数概率,但我们尚未测量它们的开销。
将得分中心化与重要性采样结合:得分中心化减去一个等于预期得分的校正项。诸如TIS或MIS之类的重要性采样方法通过使用裁剪/掩码的重要性比率重新加权得分来工作。因此,当将得分中心化组合到重要性采样之上时,我们计算的是加权得分的期望,而不是原始得分的期望。我们可以将其视为IS首先重新加权得分,然后得分中心化应用于重新加权的得分之上。
组合后的加权得分期望:令 $r_v = p_v / q_v$ 表示重要性比率,$w_v = f(r_v)$ 表示被组合的IS方法分配的权重,例如对于TIS有 $f(r) = \min(r, 2)$。得分中心化减去加权得分的期望:
就像基础的得分中心化一样,在恒定奖励下的预期更新恰好为零——通过构造,减去预期的加权得分 $\mathbb{E}_q[w_{y_t} s_{y_t}]$ 可以抵消漂移,无论加权函数 $f(r)$ 是什么。
组合方法下的 Top-k 降维推导:公式(8)中的Top-$k$缩减在这里同样适用。在建模的尾部上,采样器概率仅通过 $\hat{q}$ 已知,因此权重也是在 $\hat{q}$ 下计算的:$\hat{w}_v = f(p_v / \hat{q}_v)$,在头部它等于 $w_v$。由于 $p_v / \hat{q}_v = p_v / (\rho p_v) = 1 / \rho$ 在尾部是恒定的,中心化项保持相同的仅限头部的形式,在 $\rho$ 的位置使用标量 $\alpha$:
基础的得分中心化只是 $f = 1$ 的特例,此时给出 $\alpha = \rho$。作为完整性检查,朴素重要性采样 $f(r) = r$ 给出 $\alpha = 1$ 并且 $q_v w_v = p_v$,因此中心化项消失——精确的重要性采样已经具有零预期加权得分,所以没有什么剩下需要去中心化的了。
标量损失函数的广义表达:得分中心化可以通过将其表达在一个标量损失函数内部,在自动求导框架中高效实现。由于预期得分是逐词元得分的加权和,$\bar{s} = \sum_v q_v \nabla_\theta \log p_v$,我们可以通过使用分离的(Detached)采样器概率对训练器的对数概率进行加权,将得分中心化表达为损失。对于策略梯度,这变为:
其中 $\mathrm{sg}[\cdot]$ 表示停止梯度,即我们不对采样概率进行微分。对其微分可以恢复中心化更新:$- \nabla_\theta L = R \left( s_{y_t} - \bar{s} \right) = R \tilde{s}_{y_t}$。广义的Top-$k$版本,组合在任意重要性权重之上,遵循相同的模式,仅需要对头部词元进行计算:
广义实现代码:下面的JAX代码实现了公式(12),其中 weight_fn 指定了 $f(r)$(默认:基础SC)。与4.1节一样,train_logp 跨越整个词汇表,samp_logp 包含采样器的Top-$k$对数概率。额外的标量 samp_token_logp 是采样器对于采样词元的对数概率,即使它落在头部之外。为了数值稳定性,尾部质量的下限被设定为 eps。
import jax.numpy as jnp
from jax.lax import stop_gradient
def score_centering_loss(train_logp, samp_logp, topk_ids, sampled_token, samp_token_logp, advantage, weight_fn=lambda r: 1.0, eps=1e-6):
train_head_logp = train_logp[topk_ids]
head_weights = weight_fn(jnp.exp(train_head_logp - samp_logp))
train_tail_mass = jnp.maximum(1 - jnp.exp(train_head_logp).sum(), eps)
samp_tail_mass = jnp.maximum(1 - jnp.exp(samp_logp).sum(), eps)
tail_mass_ratio = samp_tail_mass / train_tail_mass
tail_scale = tail_mass_ratio * weight_fn(1 / tail_mass_ratio)
head_prob_residual = jnp.exp(samp_logp) * head_weights
head_prob_residual -= tail_scale * jnp.exp(train_head_logp)
logp_correction = (stop_gradient(head_prob_residual) * train_head_logp).sum()
sampled_ratio = jnp.exp(train_logp[sampled_token] - samp_token_logp)
weighted_logp = stop_gradient(weight_fn(sampled_ratio)) * train_logp[sampled_token]
return -advantage * (weighted_logp - logp_correction)
tis_weight = lambda r: jnp.minimum(r, 2.0)
mis_weight = lambda r: jnp.where((r >= 0.5) & (r <= 5.0), r, 0.0)
传递 weight_fn=tis_weight 或 weight_fn=mis_weight 以分别将得分中心化与TIS或MIS组合。正如公式(12)所示,采样的重要性权重和中心化系数都是被分离的(Detached)。
Top-k 消融实验分析:在图5中,我们在所有的实验设置中比较了Top-$k$得分中心化与完整的得分中心化。在每个设置中,$k=32$ 和 $k=128$ 的表现都与完整的得分中心化相当。公式(7)的尾部模型只需考虑Top-$k$头部之外的采样器概率质量,我们在每一步都会记录这种质量。当 $k=128$ 时,除了 30B 模型的 $\mathrm{INT8 \ W/A + INT4 \ KV}$ 设置外,头部在每个设置中平均覆盖了采样器超过 $99.9\%$ 的质量。在那个受量化扭曲最严重的尾部设置中,它平均覆盖了 $99.45\%$,在最差的批次中覆盖了 $95.8\%$——而图5显示,即使在这个设置下,$k=32$ 和 $k=128$ 仍然与完整的得分中心化相匹配。
训练目标函数的公式表达:除非另有说明,我们使用带有组内中心化奖励的REINFORCE算法:
基线校正方法的配置与细节:我们比较的所有基于IS的方法都落在一个简单的网格上(表1)。它们的区别在于统计量是按词元还是按序列计算的,重要性比率在哪里被裁剪或掩码,以及校正哪些优势符号。DPPO 使用基于二元总变差的掩码区域而不是重要性比率,但其目标仍然使用IS权重。我们使用原始论文或 verl 框架中的默认参数,并未进一步调整。所有重要性比率均针对采样器记录的概率进行计算。GSPO 使用词元比率的几何平均值,正如原论文所述。我们在同一目标函数之上应用每种校正,而不是复制每篇论文的完整方案。PG使用不带任何校正的公式(13)。
计算资源的详细分配:在所有图表中,阴影区域表示 $\pm 1$ 标准误差。图2和图3报告了保持集(Heldout)评估提示词的准确率,而图1、图4和图5报告了训练准确率(即采样轨迹的平均奖励,通过训练步数的中心移动平均线进行平滑)。表2列出了每种方法的种子数量以及重现每个图表所需的大致计算量。对于0.6B和1.7B模型,我们在每个实验中针对每种校正方法运行了3个种子。对于更昂贵的30B模型,我们最初对每种方法运行了单一种子,然后只为表现最好的方法添加了更多种子。总计算资源消耗约为 6,180 个 H100 GPU 小时。
在方法细节(Section 4)中,涉及了以下相关工作的引用说明: