Search papers, labs, and topics across Lattice.
This paper introduces Trust Region On-Policy Distillation (TrOPD), a novel approach to stabilize On-Policy Distillation (OPD) by ensuring reliable supervision from teacher models during training. By implementing trust-region constraints and outlier estimation techniques, TrOPD effectively mitigates the optimization challenges associated with distribution mismatches between teacher and student models. Experimental results demonstrate that TrOPD significantly outperforms state-of-the-art OPD methods across various tasks, including mathematical reasoning and code generation.
TrOPD stabilizes on-policy distillation by ensuring reliable teacher supervision, leading to consistent performance improvements over existing methods.
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.
