Distillation on-policy par région de confiance

Résumé

La distillation sur politique (On-Policy Distillation, OPD) est une technique fondamentale pour le post-entraînement efficace des grands modèles de langage (LLMs), avec des applications étendues dans l'apprentissage d'agents, l'amélioration multi-tâches et la compression de modèles. Cependant, l'entraînement OPD devient instable lorsque les distributions de l'enseignant et de l'élève divergent considérablement, car la supervision de l'enseignant sur les tokens générés par l'élève peut produire des gradients de politique non fiables et même provoquer un échec d'optimisation. Ce travail aborde la supervision fiable au niveau des tokens sur politique via des stratégies d'attribution de crédit, et propose la distillation sur politique avec région de confiance, TrOPD. Elle présente les caractéristiques suivantes : 1) Apprentissage sur politique avec région de confiance : TrOPD effectue l'OPD uniquement dans les régions où l'enseignant fournit une supervision fiable, atténuant la difficulté d'optimisation de l'estimateur K1 de KL inverse en cas de décalage de distribution. 2) Estimation des outliers : Pour les régions aberrantes, nous explorons l'écrêtage des gradients, le masquage et l'estimation par KL directe afin de réduire les effets négatifs d'une supervision non fiable. 3) Guidage hors politique : L'élève poursuit la génération à partir des préfixes de l'enseignant et utilise la KL directe pour imiter le guidage hors politique, encourageant l'exploration sur politique vers des régions fiables. Les expériences montrent que TrOPD surpasse systématiquement les bases OPD de l'état de l'art, notamment OPD, EOPD et REOPOLD, dans les benchmarks de raisonnement mathématique, de génération de code et de domaine général.

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.