arXiv:2605.07865cs.LGcs.AI2026-05被引 14

用强化学习技巧解决大模型训练中的不稳定问题

KL for a KL: On-Policy Distillation with Control Variate Baseline

论文配图:KL for a KL: On-Policy Distillation with Control Variate Baseline
图 1 · 摘自论文原文
  • 将推理训练转化为强化学习,引入可直接计算的基线函数
  • 在数学与科学推理任务中性能超越原始方法,接近昂贵基准
  • 无需额外计算,适合高效部署的大模型优化

On-Policy Distillation(OPD)已成为大语言模型推理领域主流的后训练范式。然而,由于单样本蒙特卡洛估计器梯度方差过高,实际训练仍不稳定,现有稳定化方法尚不成熟。本文提出vOPD(带控制变量基线的OPD),将OPD建模为策略梯度强化学习,并引入来自强化学习文献的控制变量基线——即价值函数,实现稳定。我们证明,OPD的价值函数可解析表示为学生与教师模型间逐标记的负反向KL散度,仅需已有前向传播即可获得,无需额外评论器或推理。现有方法要么需全词汇表计算逐标记反向KL,开销巨大;要么仅限于top-k支持集,引入偏差。vOPD则保持轻量级单样本估计器,通过减去分离的基线,在保证梯度无偏的同时降低方差。进一步地,对基线采用top-k近似,可进一步降低计算成本且不损害性能。在数学与科学推理基准上,vOPD持续优于原始OPD,达到最昂贵全词汇基线水平,通过严谨的强化学习方差缩减,实现了高效的OPD稳定化。

原文摘要 · Abstract (English)

On-Policy Distillation (OPD) has emerged as a dominant post-training paradigm for large language models, especially for reasoning domains. However, OPD remains unstable in practice due to the high gradient variance of its single-sample Monte Carlo estimator, and recipes for stable training are still immature. We propose vOPD (On-Policy Distillation with a control variate baseline), which casts OPD as policy-gradient RL and stabilizes it by introducing a control variate baseline-canonically a value function -- from the RL literature. We show that the OPD value function admits a closed form as the per-token negative reverse KL divergence between the student and the teacher, available directly from the already-computed forward pass with no additional critic or inference. Existing stabilization methods either compute the full token-level reverse KL over the entire vocabulary, adding significant overhead, or restrict it to a top-k support, biasing the objective. vOPD instead preserves the lightweight single-sample estimator, subtracting the value function as a detached baseline to keep the gradient unbiased while reducing variance. Furthermore, we show that a top-k approximation of the baseline further lowers cost without compromising performance. Across mathematical and scientific reasoning benchmarks, vOPD consistently outperforms vanilla OPD and matches the most expensive full-vocabulary baseline, offering an efficient stabilization of On-Policy Distillation through principled RL variance reduction.

大模型训练强化学习知识蒸馏稳定性优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。