Unmasking On-Policy Distillation: Where It Helps, Where It Hurts, and Why
On-policy distillation provides dense per-token supervision for reasoning models, but its effectiveness varies. This research introduces a training-free diagnostic framework that quantifies the alignment between distillation signals and the ideal gradient at the per-token level. Findings show that distillation guidance aligns better on incorrect rollouts, and the optimal context depends on student capacity and task, with no universal best configuration.
content type paperpublished July 2026
Unmasking On-Policy Distillation: Where It Helps, Where It Hurts, and Why
AuthorsMohammadreza Armandpour*, Fatih Ilhan*, David Harrison, Ajay Jaiswal, Duc N.M Hoang, Fartash Faghri, Yizhe Zhang, Minsik Cho, Mehrdad Farajtabar
View publication
On-policy distillation offers dense, per-token supervision for training reasoning models; however, it remains unclear under which conditions this signal is beneficial and under which it is detrimental. Which teacher model should be used, and in the case of self-distillation, which specific context should serve as the supervisory signal? Does the optimal choice vary from one token to the next? At present, addressing these questions typically requires costly training runs whose aggregate performance metrics obscure the dynamics at the level of individual tokens. We introduce a training-free diagnostic framework that operates at the highest resolution: per token, per question, and per teacher. We derive an ideal per-node gradient defined as the parameter update that maximally increases the student’s probability of success. We then develop a scalable targeted-rollout algorithm to estimate this gradient efficiently, even for long chains of intermediate thoughts. The gradient alignment score, defined as the cosine similarity between this ideal gradient and any given distillation gradient, quantifies the extent to which a particular configuration approximates the ideal signal. Across a range of self-distillation settings and external teacher models, we observe that distillation guidance exhibits substantially higher alignment with the ideal on incorrect rollouts than on correct ones, where the student already performs well and the teacher’s signal tends to become noisy. Furthermore, we find that the optimal distillation context depends jointly on the student model’s capacity and the target task, and that no single universally effective configuration emerges. These findings motivate the use of per-task, per-token diagnostic analyses for distillation.
- Equal contribution
Distillation Scaling Laws
July 1, 2025research area Methods and Algorithms, research area Speech and Natural Language Processingconference ICML
We propose a distillation scaling law that estimates distilled model performance based on a compute budget and its allocation between the student and teacher. Our findings mitigate the risks associated with large-scale distillation by enabling compute-optimal allocation for both the teacher and student to maximize student performance. We provide compute-optimal distillation recipes for two key scenarios: when a teacher already exists, and when a…
Read more
Extracurricular Learning: Knowledge Transfer Beyond Empirical Distribution
June 10, 2021research area Computer VisionWorkshop at CVPR
Knowledge distillation has been used to transfer knowledge learned by a sophisticated model (teacher) to a simpler model (student). This technique is widely used to compress model complexity. However, in most applications the compressed student model suffers from an accuracy gap with its teacher. We propose extracurricular learning, a novel knowledge distillation method, that bridges this gap by (1) modeling student and teacher output…
Read more