Trust Region On-Policy Distillation

Zusammenfassung

On-Policy Distillation (OPD) ist eine grundlegende Technik für effizientes Post-Training großer Sprachmodelle (Large Language Models, LLMs) mit breiten Anwendungen im Agentenlernen, in der Mehrfachaufgabenverbesserung und in der Modellkompression. Allerdings wird das OPD-Training instabil, wenn die Verteilungen von Lehrer und Schüler erheblich voneinander abweichen, da die Lehreraufsicht über von Schülern generierte Token zu unzuverlässigen Policy-Gradienten und sogar zu Fehlern bei der Optimierung führen kann. Diese Arbeit befasst sich mit zuverlässiger On-Policy-Token-Überwachung durch Kreditzuweisungsstrategien und schlägt Trust Region On-Policy Distillation (TrOPD) vor. Es zeichnet sich durch folgende Merkmale aus: 1) Trust-Region-On-Policy-Lernen: TrOPD führt OPD nur in Bereichen durch, in denen der Lehrer zuverlässige Aufsicht bietet, und mildert so die Optimierungsschwierigkeiten des K1-Reverse-KL-Schätzers bei Verteilungsdiskrepanz. 2) Ausreißerschätzung: Für Ausreißerbereiche untersuchen wir Gradienten-Clipping, Maskierung und Forward-KL-Schätzung, um die negativen Auswirkungen unzuverlässiger Aufsicht zu verringern. 3) Off-Policy-Anleitung: Der Schüler setzt die Generierung ab Lehrer-Präfixen fort und verwendet die Forward-KL, um die Off-Policy-Anleitung zu imitieren, was die On-Policy-Exploration in Richtung zuverlässiger Regionen fördert. Experimente zeigen, dass TrOPD durchgängig bessere Ergebnisse erzielt als die aktuellen State-of-the-Art-OPD-Baselines, darunter OPD, EOPD und REOPOLD, bei Benchmarks zum mathematischen Denken, zur Codegenerierung und in allgemeinen Domänen.

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.