Leren voorbij de leraar: Gegeneraliseerde on-policy distillatie met beloningsextrapolatie

Learning beyond Teacher: Generalized On-Policy Distillation with Reward Extrapolation

February 12, 2026

Auteurs: Wenkai Yang, Weijie Liu, Ruobing Xie, Kai Yang, Saiyong Yang, Yankai Lin

cs.AI

Samenvatting

On-policy distillatie (OPD), waarbij de student wordt afgestemd op de logit-verdeling van de leraar op door de student gegenereerde trajecten, heeft sterke empirische verbeteringen aangetoond in de prestaties van de student en overtreft vaak off-policy distillatie en reinforcement learning (RL) paradigma's. In dit werk tonen we eerst theoretisch aan dat OPD een speciaal geval is van dense RL met een KL-beperking, waarbij de beloningsfunctie en de KL-regularisatie altijd even zwaar worden gewogen en het referentiemodel elk willekeurig model kan zijn. Vervolgens stellen we het Generalized On-Policy Distillation (G-OPD) raamwerk voor, dat het standaard OPD-doel uitbreidt door een flexibel referentiemodel en een schaalfactor voor de beloning te introduceren die het relatieve gewicht van de beloningsterm ten opzichte van de KL-regularisatie regelt. Door middel van uitgebreide experimenten met wiskundige redeneer- en codegeneratietaken leiden we twee nieuwe inzichten af: (1) Het instellen van de schaalfactor voor de beloning op groter dan 1 (d.w.z. beloningsextrapolatie), wat we ExOPD noemen, levert consistente verbetering op ten opzichte van standaard OPD over een reeks van leraar-student grootte-combinaties. Met name in de setting waar we kennis van verschillende domeinexperts, verkregen door domeinspecifieke RL op hetzelfde studentmodel toe te passen, terug samenvoegen in de oorspronkelijke student, stelt ExOPD de student in staat om zelfs de prestatiegrens van de leraar te overtreffen en de domeinleraren te verslaan. (2) Voortbouwend op ExOPD, ontdekken we verder dat in de strong-to-weak distillatie-setting (d.w.z. het distilleren van een kleinere student uit een grotere leraar), het uitvoeren van beloningscorrectie door het referentiemodel te kiezen als het basis model van de leraar vóór RL een nauwkeuriger beloningssignaal oplevert en de distillatieprestatie verder verbetert. Deze keuze vereist echter toegang tot de variant van de leraar vóór RL en brengt meer rekenkosten met zich mee. We hopen dat ons werk nieuwe inzichten biedt voor toekomstig onderzoek naar OPD.

English

On-policy distillation (OPD), which aligns the student with the teacher's logit distribution on student-generated trajectories, has demonstrated strong empirical gains in improving student performance and often outperforms off-policy distillation and reinforcement learning (RL) paradigms. In this work, we first theoretically show that OPD is a special case of dense KL-constrained RL where the reward function and the KL regularization are always weighted equally and the reference model can by any model. Then, we propose the Generalized On-Policy Distillation (G-OPD) framework, which extends the standard OPD objective by introducing a flexible reference model and a reward scaling factor that controls the relative weight of the reward term against the KL regularization. Through comprehensive experiments on math reasoning and code generation tasks, we derive two novel insights: (1) Setting the reward scaling factor to be greater than 1 (i.e., reward extrapolation), which we term ExOPD, consistently improves over standard OPD across a range of teacher-student size pairings. In particular, in the setting where we merge the knowledge from different domain experts, obtained by applying domain-specific RL to the same student model, back into the original student, ExOPD enables the student to even surpass the teacher's performance boundary and outperform the domain teachers. (2) Building on ExOPD, we further find that in the strong-to-weak distillation setting (i.e., distilling a smaller student from a larger teacher), performing reward correction by choosing the reference model as the teacher's base model before RL yields a more accurate reward signal and further improves distillation performance. However, this choice assumes access to the teacher's pre-RL variant and incurs more computational overhead. We hope our work offers new insights for future research on OPD.