On-policy distillation sometimes fails due to weak learning signals
Why On-Policy Distillation Sometimes Fails: Vanishing Learning Signals
Machine LearningComputation and Language
Summary
Sometimes when copying what a smarter AI model does to teach a smaller AI model, the smaller one stops learning early and gets stuck. The authors studied this problem and found that the smaller model’s learning signals, which tell it how to improve, become very weak before it has learned enough. They also found that when the teacher and student models are very similar, learning works better, and that the student’s knowledge changes very little during this process, which might be why learning stops. This research helps explain why some AI teaching methods don’t work as expected and suggests learning signals and model similarity are important factors.
What this means in practice
- •For ai engineers: Improve model training by monitoring learning signals and maintaining similarity between teacher and student models to avoid early learning plateaus.
- •For software developers: Develop more reliable AI-assisted coding and reasoning tools by addressing weak learning signals during model distillation.
Authors
Lei Zhao, Qichao Zhao, Bowen Zuo, Qishi Zhan
Abstract
On-policy distillation (OPD) enables effective capability transfer between language models, yet the mechanisms underlying its failures are not fully understood. Across code generation and mathematical reasoning, OPD with larger-scale teachers exhibits early loss plateaus, with an average final loss reduction of 25.1% after 200 updates, compared with 96.2% for self-RL teachers, obtained by further reinforcement learning (RL) training of the initial student. To understand this difference, we analyze OPD as an idealized continuous-time dynamical system in the small-learning-rate limit. Our training-log diagnostics associate these plateaus with an early decline in a gradient-based learning-signal proxy while substantial loss remains; these measurements do not establish why the underlying gradient weakens. We further prove a local recovery guarantee for teachers sufficiently close to the initial student in a shared parameterization under regularity conditions, offering a conditional explanation for the success of self-RL teachers in our experiments. Across runs with and without loss plateaus, we observe small relative parameter changes (0.025-0.098%) and high similarity between the student's representations before and after OPD (linear CKA $>0.98$ across layers). These observations suggest that limited representation adaptation may contribute to learning-signal collapse, a hypothesis that remains to be tested. Code is available at https://github.com/leizhao7/opd-learning-signals.