Source-linked AI summary
Learning beyond Teacher: Generalized On-Policy Distillation with Reward Extrapolation
Wenkai Yang, Weijie Liu, Ruobing Xie, Kai Yang, Saiyong Yang, Yankai Lin
TL;DR
The paper studies the limited mechanistic understanding of on-policy distillation and develops a generalized formulation connected to dense KL-constrained RL. G-OPD introduces flexible reference models and reward scaling, while experiments show that reward extrapolation improves distillation and can surpass domain teachers. In strong-to-weak distillation, reward correction further improves performance but requires the teacher’s pre-RL model and additional computation.
Problem
OPD has shown empirical effectiveness, but its mechanistic understanding and full potential remain limited.
Method
The paper derives OPD as a special case of dense KL-constrained RL and proposes G-OPD with flexible reference models and reward scaling.
Results
ExOPD consistently improves over OPD across teacher-student settings, can surpass domain teachers in multi-teacher distillation, and improves strong-to-weak distillation with reward correction.
Takeaways & Limitations
Reward extrapolation provides a route for students to learn beyond teacher capability boundaries, including when merging multiple domain experts.
Takeaways & Limitations
Using the teacher’s pre-RL model for reward correction assumes access to that model and incurs additional computational overhead; sparse RL rewards can also make optimization inefficient.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper addresses limited mechanistic understanding of on-policy distillation by connecting OPD to dense KL-constrained reinforcement learning and generalizing its objective. Experiments show that reward extrapolation can surpass teacher capability boundaries, including in multi-teacher distillation.
- Contributions: OPD is theoretically characterized as a special case of dense KL-constrained RL with equally weighted token rewards and KL regularization.The reference model in this formulation can be chosen arbitrarily.
- Contributions: G-OPD extends standard OPD with a flexible reference model and a reward scaling factor controlling reward weight relative to KL regularization.The framework is designed to analyze how these choices affect distillation across settings.
- Key findings: Reward interpolation produces students with performance and response length between the reference and teacher, whereas reward extrapolation can surpass the teacher’s capability boundary.The paper names the reward-extrapolation variant ExOPD.
- Key findings: ExOPD extends to multi-teacher distillation, enabling a unified student to surpass all domain teachers.The experiments cover math reasoning and code generation domains.
2 Related Work
Prior knowledge-distillation work has largely used teacher-generated trajectories, while on-policy distillation instead trains on student-generated trajectories and aligns token-level logits with the teacher. Empirically, OPD has been reported to make distillation faster and more effective than off-policy approaches.
- Off-Policy Distillation: Off-policy distillation typically trains students on teacher-generated trajectories using token-logit KL alignment or related objectives.
- On-Policy Distillation: OPD samples trajectories from the student and aligns the student’s token-level logit distribution with the teacher’s.
- On-Policy Distillation: OPD realizes dense on-policy learning by applying this alignment across tokens on student-generated trajectories.
- On-Policy Distillation: Empirical studies report that OPD achieves faster and more effective distillation than off-policy distillation.
3 Methodology
This section formulates OPD as dense KL-constrained reinforcement learning and generalizes it with a flexible reference model and adjustable reward weighting. It analyzes reward interpolation and extrapolation, including reference-model correction for strong-to-weak distillation.
- 3.1 Preliminaries: OPD trains on student-generated trajectories, combines on-policy learning with dense credit assignment, and minimizes reverse KL divergence between student and teacher distributions.Unlike off-policy distillation, the student learns from supervision on trajectories generated by its own policy.
- 3.2 Generalized On-Policy Distillation: OPD is equivalent to a KL-constrained RL objective whose reward is the teacher-to-reference log-probability shift, with reward and KL terms weighted equally.The reference model can be chosen flexibly, including models of different sizes from the teacher.
- 3.2 Generalized On-Policy Distillation: G-OPD introduces a reference model and reward scaling factor λ, enabling control over the reward term’s weight relative to KL regularization while retaining dense credit assignment.Compared with OPD, λ removes the fixed 1:1 weighting between reward and KL terms.
- 3.2 Generalized On-Policy Distillation: For 0 < λ < 1, G-OPD interpolates between the reference model and standard OPD; for λ > 1, it extrapolates the teacher-reference log-probability shift.The λ > 1 regime is termed reward extrapolation and may produce students that outperform the teacher in certain cases.
- 3.2 Generalized On-Policy Distillation: In strong-to-weak distillation, choosing the teacher’s pre-RL base model as the reference may improve distillation by correcting the reward signal.This choice requires access to an additional teacher-base model and incurs extra computational cost.
4 Experiments and Analysis
Experiments evaluate G-OPD and ExOPD across math, code, multi-teacher, and strong-to-weak distillation settings. Reward extrapolation improves over standard OPD, can surpass domain teachers in multi-teacher merging, and produces higher rewards alongside longer responses.
- Experimental setup: Experiments cover math reasoning and code generation, including same-size single-teacher, multi-teacher, and strong-to-weak distillation settings.The study uses four math benchmarks and three code-generation benchmarks, with Qwen3-4B variants for same-size experiments.
- Single-teacher distillation: Standard OPD recovers the domain teacher’s post-training behavior, while reward interpolation places performance and response length between the base model and teacher.Performance and response length increase monotonically with λ during reward interpolation.
- Single-teacher distillation: λ = 1.25 consistently outperforms OPD and the domain teacher across single-teacher benchmarks, whereas λ = 1.5 may destabilize training and degrade performance.The paper attributes excessive extrapolation risk to hacking the implicit reward and length bias.
- Robustness analysis: Additional teacher RL training does not eliminate ExOPD’s improvement, as the teacher’s continued-RL gains are smaller than ExOPD’s gains with fewer steps.The comparison uses an additional 100 RL steps and also evaluates teachers trained for 1200 steps in Appendix C.
- Multi-teacher distillation: ExOPD is the only evaluated multi-teacher method that produces a unified student surpassing both domain teachers on all benchmarks.It consistently outperforms OPD, while SFT is sub-optimal and ExPO lacks consistent controllability.
- Multi-teacher distillation: ExOPD achieves higher training rewards than OPD but produces longer responses and higher response entropy in multi-teacher training.The higher entropy is attributed to longer responses increasing response diversity.
A Detailed Math Derivations
The appendix derives expected gradients for OPD and G-OPD, including reduced equation forms and practical gradient approximations. It also points to tables containing training hyper-parameters for RL, G-OPD, and SFT.
- The derivations calculate the expected gradients of the OPD objective in Eq. (4).
- Equations (16) and (18) are reduced to subsequent forms in the derivation.
- Training hyper-parameters are listed separately for math and code RL, G-OPD, and SFT.
- Using a discount factor of 0, recent OPD studies approximate the gradient in practice.
- The appendix provides the corresponding approximated gradient expression for G-OPD.
B Detailed Training Settings
The appendix specifies training settings for RL, G-OPD, and SFT, including optimization-step schedules and fairness controls for comparing distillation methods.
- Math and code RL training hyper-parameters are provided in Tables 4 and 5.
- G-OPD uses 50 optimization steps for same-size teacher-student pairs and 100 steps for strong-to-weak distillation.
- Larger prompt sizes produce smoother convergence under the same prompt-size and rollout conditions in preliminary experiments.
- SFT matches the teacher-generated trajectory count to the student-generated count used in OPD and ExOPD.
- SFT keeps optimization-step counts consistent with the corresponding G-OPD experiment for fair comparison.
C Results of Distillation from Domain Teachers with Sufficient RL Trainings
Additional experiments train domain teachers with sufficient RL and evaluate ExOPD under the same settings as Section 4.1, supporting its generalizability and effectiveness in this case.
- With domain teachers trained for 1200 RL steps, Table 8 evaluates ExOPD using the Section 4.1 experimental settings.
- The Table 8 results demonstrate ExOPD’s generalizability and effectiveness under stronger teachers with sufficient RL training.
D Prompt Templates
The experiments use separate prompt templates for math reasoning and code generation, specifying the requested reasoning or programming output format.
- The paper provides prompt templates for both training and evaluation at the end of the appendix.
- Training and Evaluation Prompt Template for Math Reasoning: The math template asks the model to reason step by step and place the final answer within boxed formatting.
- Training and Evaluation Prompt Template for Code Generation: The code template asks the model to write Python code, present it in a Python code block, and think before writing.