Apprendimento oltre l'Insegnante: Distillazione Generalizzata On-Policy con Estrapolazione della Ricompensa

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

February 12, 2026

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

cs.AI

Abstract

La distillazione on-policy (OPD), che allinea lo studente con la distribuzione dei logit del teacher su traiettorie generate dallo studente, ha dimostrato significativi vantaggi empirici nel migliorare le prestazioni dello studente, superando spesso i paradigmi della distillazione off-policy e dell'apprendimento per rinforzo (RL). In questo lavoro, dimostriamo prima teoricamente che l'OPD è un caso speciale di RL con vincolo KL denso, in cui la funzione di ricompensa e la regolarizzazione KL sono sempre ponderate in modo uguale e il modello di riferimento può essere qualsiasi modello. Successivamente, proponiamo il framework di Distillazione On-Policy Generalizzata (G-OPD), che estende l'obiettivo OPD standard introducendo un modello di riferimento flessibile e un fattore di scala della ricompensa che controlla il peso relativo del termine di ricompensa rispetto alla regolarizzazione KL. Attraverso esperimenti completi su compiti di ragionamento matematico e generazione di codice, deriviamo due nuove intuizioni: (1) Impostare il fattore di scala della ricompensa maggiore di 1 (cioè, l'estrapolazione della ricompensa), che chiamiamo ExOPD, migliora costantemente rispetto all'OPD standard in una serie di accoppiamenti di dimensione teacher-studente. In particolare, nello scenario in cui uniamo la conoscenza di diversi esperti di dominio, ottenuta applicando RL specifico per dominio allo stesso modello studente, nuovamente nello studente originale, ExOPD consente allo studente di superare persino il confine prestazionale del teacher e di ottenere risultati migliori dei teacher di dominio. (2) Basandoci su ExOPD, scopriamo inoltre che, nello scenario di distillazione da forte a debole (cioè distillare uno studente più piccolo da un teacher più grande), effettuare una correzione della ricompensa scegliendo come modello di riferimento il modello base del teacher prima del RL fornisce un segnale di ricompensa più accurato e migliora ulteriormente le prestazioni della distillazione. Tuttavia, questa scelta presuppone l'accesso alla variante pre-RL del teacher e comporta un sovraccarico computazionale maggiore. Speriamo che il nostro lavoro offra nuove prospettive per la ricerca futura sull'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.