解决大模型蒸馏中教师学生分布差异导致的训练不稳问题
Trust Region On-Policy Distillation
- 仅在教师监督可靠区域进行蒸馏，避免错误梯度干扰
- 通过裁剪、掩码和前向KL估计降低异常区域影响
- 用教师前缀引导学生探索可靠生成区域，适合模型压缩与强化学习
在线策略蒸馏（OPD）是高效微调大语言模型的核心技术，广泛应用于智能体学习、多任务增强与模型压缩。然而当教师与学生分布差异显著时，教师对学生产生令牌的监督可能导致不可靠策略梯度，引发优化失败。本文提出可信区域在线策略蒸馏（TrOPD），核心包括：1）仅在教师提供可靠监督的区域执行OPD，缓解分布不匹配下逆KL估计器的优化困难；2）对异常区域采用梯度裁剪、掩码与前向KL估计，降低不可靠监督的负面影响；3）学生沿用教师前缀生成，并以正向KL模仿离线指导，推动在线探索向可靠区域转移。实验表明，TrOPD在数学推理、代码生成与通用基准上均优于现有最优的OPD基线方法，包括OPD、EOPD与REOPOLD。
原文摘要 · Abstract (English)
On-Policy Distillation (OPD) is a fundamental technique for efficient post-training of large language models (LLMs), with broad applications in agent learning, multi-task enhancement, and model compression. However, OPD training becomes unstable when the teacher and student distributions differ substantially, as teacher supervision on student-generated tokens may yield unreliable policy gradients and even cause optimization failure. This work addresses reliable on-policy token-level supervision through credit assignment strategies, and proposes Trust Region On-Policy Distillation, TrOPD. It features the following characteristics: 1) Trust-Region On-Policy Learning: TrOPD performs OPD only in regions where the teacher provides reliable supervision, mitigating the optimization difficulty of the K1 reverse-KL estimator under distribution mismatch. 2) Outlier Estimation: For outlier regions, we explore gradient clipping, masking, and forward-KL estimation to reduce the adverse effects of unreliable supervision. 3) Off-Policy Guidance: The student continues generation from teacher prefixes and uses forward KL to imitate off-policy guidance, encouraging on-policy exploration toward reliable regions. Experiments show that TrOPD consistently outperforms SoTA OPD baselines, including OPD, EOPD, and REOPOLD, across mathematical reasoning, code generation, and general-domain benchmarks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。