只用生成结果的前缀做教师监督,训练速度提升47倍
Fast and Effective On-policy Distillation from Reasoning Prefixes
- 仅对学生模型输出的前缀进行教师监督,减少采样开销
- 在数学和跨域任务上性能接近完整策略蒸馏,训练耗能降低2至47倍
- 适合需要高效微调大模型的场景,尤其长文本生成任务
在线策略蒸馏(OPD)通过从学生模型采样轨迹,并以教师模型在词元层面提供监督,避免了仅依赖可验证终端奖励的问题,相比离线策略蒸馏具有更好的泛化能力。然而,OPD 在训练过程中需实时采样学生策略,导致训练成本高昂,尤其在生成长序列时更为明显。我们初步分析发现,在 OPD 中,训练信号往往集中在输出前缀部分,即使短长度的教师生成前缀也能显著帮助学生生成正确答案。基于此,我们提出一种简单而有效的改进:仅对学生生成输出的前缀应用蒸馏目标,并在蒸馏过程中提前终止采样。在一系列 AI-for-Math 和跨域基准测试中,该方法在保持与完整 OPD 相当性能的同时,将训练浮点运算量(FLOP)降低了2至47倍。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。