On-policy distillation boosts model scale transfer in reinforcement learning
Scaling Properties of Same-Family On-Policy Distillation
Machine Learning
Summary
It’s hard to know how well lessons learned by smaller AI models help bigger ones. The paper looks at a way to teach models, called on-policy distillation, and finds that smaller expert models can actually help bigger ones learn better. The improvement depends on both the size of the models and how skilled the teacher model is. The researchers also found that bigger teachers don’t always work better when teaching smaller students.
What this means in practice
- •For machine learning engineers: Improve training pipelines by optimizing model size combinations for reinforcement learning distillation setups to boost student model performance.
- •For ai system architects: Design efficient reinforcement learning systems by selecting appropriately sized teacher models to maximize knowledge transfer to student models.
Authors
Yuntai Bao, Qinfeng Li, Guoqing Jiang, Liwei Chen, Zhiheng Qin, Xuanping Li, Wenqi Zhang, Xuhong Zhang
Abstract
*Reinforcement learning (RL)* can induce substantial reasoning capabilities in large language models (LLMs), but how much of this capability transfers across model scales, and how quickly, remains unclear. We study the scaling properties of *on-policy distillation (OPD)* across *weak-to-strong*, *same-base*, and *strong-to-weak* teacher--student setups. We find that early OPD training dynamics uniformly exhibit a regular *useful-transfer* regime, in which held-out accuracy (the *gold score*, $G$) rises approximately linearly in $d=\sqrt{\mathrm{KL}(π_θ\Vert π_{\mathrm{ref}})}$, the square root of token-level reverse KL divergence from the student initialization. In every observed weak-to-strong pair, the student's peak gold score exceeds its teacher's own, so a compact RL expert can transfer capability to a much larger student via OPD. To estimate OPD outcomes, we fit *power laws* for how $G_{\mathrm{peak}}$ and the slope of the useful-transfer regime scale with student and teacher parameter counts and with teacher gold score. These laws show that peak gold score improves with teacher scale only up to roughly the student's scale, and that at a matched gold score smaller teachers transfer better, so a teacher's score alone does not define its supervision value. We also study the scaling effects of two OPD variants, bootstrapping weak-to-strong OPD, and the degree of on-policy supervision.