Rethinking Training-Inference Mismatch in LLM Reinforcement Learning: Where It Arises and How to Correct It
Tianrun Yu, Kaixiang Zhao, Shangzhe Li, Yuxiao Yang, Porter Jenkins, Weitong Zhang, Taylor W. Killian
cs.LG, cs.AI
2026-09-26
CIS把训练推理失配写成与置信度近似无关的log-odds位移并只截上尾，三款MoE在五个数学基准上平均全最高，Qwen1.5-MoE从无校正30.99升到34.78。
当前 LLM 的可验证奖励强化学习(RLVR)把 rollout 和梯度拆给两套引擎:推理侧(vLLM、SGLang)采样,训练侧(FSDP、Megatron)反传。参数相同,kernel、归约顺序和数值精度不同,同一个 token 的概率对不上。名义 on-policy,实际已经 off-policy。MoE 上更狠,因为 top-k 路由是离散决策,一点点扰动就能换专家,训练会不稳甚至崩掉。
常规补救是重要性采样,用训练侧概率 p 除以推理侧概率 q 得到 k=p/q,再乘到 GRPO 的 token 损失上。两边偶尔吵得很凶,k 会大到梯度方差失控。TIS 给 k 加固定上帽,IcePop 把落在区间外的 token 掩掉,KPop 按二元 KL 掩码,还有序列级截断和整条管线改 FP16。这些规则都直接打在 k 或两边概率上。k 却把两件事混在一起:引擎偏差,以及这个 token 有多自信。同样的偏差,低置信处把 k 甩出很远,高置信处几乎不动。固定阈值于是主要砸在低置信 token 上,截断偏差也堆在那里。
失配可以精确写成 k = p + (1-p)exp(ε),其中 ε 是两边给该 token 的 log-odds 之差,由 softmax 前的逐维 logit 扰动决定。稠密模型上扰动像 bf16 舍入,几乎贴零;MoE 主体也贴零,但拖一条重尾。24 层里平均 3.69 层的 top-4 专家集合不一致,只有 8.0% 的位置完全一致;不一致层从 0 增到至少 7,|δ|>0.5 的比例从 4.51% 升到 18.48%。即便当前步路由完全一致,MoE 扰动标准差仍约是稠密对照的三倍。
ε 几乎不随置信度变:六个数量级的不确定度上中位数贴零,离散程度变化不到两倍,Spearman 相关 -0.028。方差是单侧的,训练后 MoE 上 k>1 的 token 贡献了 Var(k) 的 99.96%。精确校正的 E[e^{2ε}] 有 9.2×10^6,乘上 (1-p) 后二阶矩仍是 22.9,无校正时这个期望等于 1。
CIS 只截上尾:把 exp(ε) 卡在 1+λ,映射回比值就是 min(k, 1+λ(1-p))。token 越自信,允许的 k 越靠近 1。对 1-p 再垫地板 κ=5×10^{-3},否则帽子会窄过 log 概率的存储精度,平均准确率从 34.78 掉到 29.02。默认 λ=2.3。额外成本是一次逐元素计算;权重要 detach,推理侧必须用采样当时记下的生 log 概率。精确 IS 里可能炸掉的二阶矩被换成不超过 (1+λ)^2 的项,偏差由被截掉的超额位移控制。
三款 MoE,同一套 GSM8K 配方(7473 题、三 epoch、87 步、无 KL 惩罚),五个数学基准 greedy pass@1,三颗种子。CIS 在三款模型上都拿到五基准平均最高。
| 方法 | Qwen1.5-MoE | DeepSeek-V2-Lite | Qwen3-30B-A3B |
| 无校正 | 30.99 | 35.83 | 68.26 |
| Exact ratio | 32.48 | 35.47 | 68.72 |
| IcePop | 34.18 | 36.73 | 68.99 |
| TIS (C=2) | 31.40 | 36.02 | 69.14 |
| GSPO | 28.90 | 36.47 | 69.23 |
| CIS | 34.78 | 37.16 | 69.88 |
Qwen1.5-MoE 上 CIS 的 GSM8K 是 67.88,无校正是 61.21。FP16 平均 31.02,几乎等于无校正。若干迁移列相对 IcePop 的优势小于种子波动,论文不把单列差当成统计分离。Qwen3-30B 基座 GSM8K 已有 95.22,训练任务几乎没余量。
消融把三块设计钉死。两侧一起截,平均从 34.78 掉到 6.01,低于基座 28.93。固定帽 C=1.13 截掉相近比例的 token(4.3% 对 3.8%),平均只剩 8.88。去掉地板掉到 29.02。λ=0 崩到 8.75;凡是正阈值都高于无校正的 30.99,默认 2.3 最高,其余正值在 32.53 到 34.08。训练上 CIS 奖励 0.867、峰值梯度 3.0,精确比值峰值梯度 69.5。CIS 大约动 3.8% 的 token。总体截断偏差 0.0064,低于 TIS 的 0.0078;额外方差 0.0081,低于精确校正的 0.0132。
MoE 上做 RLVR,训练-推理失配会把梯度打飞,这是稳定性问题。CIS 不改数值精度、不回放路由,只改 token 损失权重,能嵌进现有 TIS/IcePop 那条分支。对已经在截断 IS 的团队,这是一次带机制解释的渐进替换:同样截大约 4% 的 token,把刀从低置信比值挪到位移坐标,held-out 才站得住。它管的是残差失配的梯度冲击,和 FP16、routing replay 互补。Qwen3-30B 相对最强基线只高 0.65 分,小模型上那 3.8 分更有意义。
若干迁移基准上相对 IcePop 的优势落在种子标准差里;域外知识与代码的池化准确率也最高,边同样在种子散布内。路由分析是观测分层不是干预,前缀期路由分歧拆不干净。ε 对置信度只是近似不变,高置信端还会略变宽。
实验面窄:RL 只在 GSM8K 上跑了 87 步,三款全是 MoE,稠密模型只出现在失配测量里,主表没有 routing replay。k 只校正给定前缀下的 next-token,不重加权前缀分布。定理假定 token 贡献独立,相关时要用有效样本量。λ 贴 0 会崩,正值区间里平均还能摆 2.3 分,换框架后默认 2.3 能不能直接搬,没测。