arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2602.15260cs.LGcs.AI
收藏

基于推理前缀的高效在线策略蒸馏

Fast and Effective On-policy Distillation from Reasoning Prefixes

Optum AI(Optum人工智能部门)

Dongxu Zhang, Zhichao Yang, Sepehr Janghorbani, Jun Han, Andrew Ressler, Qian Qian, Gregory D. Lyng, Sanjit Singh Batra, Robert E. Tillman

AI总结:

本文提出基于推理前缀的在线策略蒸馏方法,通过仅对生成输出的前缀应用蒸馏目标并提前终止采样,有效降低训练成本,同时保持与完整OPD相当的性能。

摘要:

在线策略蒸馏(OPD)通过从学生模型采样轨迹并以token级别用教师模型监督,避免仅依赖可验证的终端奖励,可以比离线策略蒸馏产生更好的泛化能力。然而,OPD在训练过程中需要昂贵的在线采样学生策略,这大大增加了训练成本,尤其是对于长响应。我们的初步分析表明,在OPD过程中,训练信号往往集中在每个输出的前缀中,且即使一个短的教师生成前缀也能显著帮助学生生成正确答案。受这些观察的启发,我们提出了一种简单而有效的OPD修改方法:我们仅将蒸馏目标应用于学生生成输出的前缀,并在蒸馏过程中提前终止每个采样。在一系列AI-for-Math和外领域基准测试中,基于前缀的在线蒸馏在性能上与完整OPD匹配,同时将训练FLOP减少2倍至47倍。

英文摘要:

On-policy distillation (OPD), which samples trajectories from the student model and supervises them with a teacher at the token level, avoids relying solely on verifiable terminal rewards and can yield better generalization than off-policy distillation. However, OPD requires expensive on-the-fly sampling of the student policy during training, which substantially increases training cost, especially for long responses. Our initial analysis shows that, during OPD, training signals are often concentrated in the prefix of each output, and that even a short teacher-generated prefix can significantly help the student produce the correct answer. Motivated by these observations, we propose a simple yet effective modification of OPD: we apply the distillation objective only to prefixes of student-generated outputs and terminate each sampling early during distillation. Experiments on a suite of AI-for-Math and out-of-domain benchmarks show that on-policy prefix distillation matches the performance of full OPD while reducing training FLOP by 2x-47x.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑