Trust Region On-Policy Distillatie

Samenvatting

On-Policy Destillatie (OPD) is een fundamentele techniek voor efficiënte nabehandeling van grote taalmodellen (LLM's) en kent brede toepassingen in agentleren, meertakenverbetering en modelcompressie. OPD-training wordt echter instabiel wanneer de verdelingen van de leraar en student aanzienlijk verschillen, omdat toezicht van de leraar op door de student gegenereerde tokens kan leiden tot onbetrouwbare beleidsgradiënten en zelfs optimalisatiefalen. Dit werk pakt betrouwbaar toezicht op token-niveau via krediettoewijzingsstrategieën aan en stelt Trust Region On-Policy Destillatie voor, TrOPD. Het kenmerkt zich door de volgende eigenschappen: 1) Trust-Region On-Policy Leren: TrOPD voert OPD alleen uit in gebieden waar de leraar betrouwbaar toezicht biedt, wat de optimalisatieproblemen van de K1 reverse-KL schatter onder distributiemismatch verlicht. 2) Uitbijterschatting: Voor uitbijtergebieden verkennen we gradiëntclipping, masking en forward-KL schatting om de nadelige effecten van onbetrouwbaar toezicht te verminderen. 3) Off-Policy Begeleiding: De student gaat verder met genereren vanaf lerar prefixen en gebruikt forward KL om off-policy begeleiding te imiteren, wat on-policy verkenning naar betrouwbare gebieden stimuleert. Experimenten tonen aan dat TrOPD consistent beter presteert dan state-of-the-art OPD-baselines, waaronder OPD, EOPD en REOPOLD, op het gebied van wiskundig redeneren, codegeneratie en algemene-domein benchmarks.

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.