New method improves learning from fewer teacher tokens in on-policy distillation
DIAL-OPD: Learning More from Fewer Tokens in On-Policy Distillation
Computation and LanguageArtificial Intelligence
Summary
Sometimes in teaching AI models, showing fewer examples or parts to learn from can actually lead to better learning than showing everything. The authors found that picking the right tokens (small pieces of information) to learn from helps the student AI learn better. They made a new way to choose these tokens based on balancing how both the teacher and student models see their importance. Their method, called DIAL-OPD, works better than many existing approaches even when using less information from the teacher.
What this means in practice
- •For machine learning engineers: Optimize training large language models by selectively supervising fewer tokens without losing accuracy, reducing computation costs.
- •For nlp system developers: Improve reasoning accuracy of natural language models by applying smarter token-level distillation in deployment pipelines.
Authors
Anhao Zhao, Haoran Xin, Junlong Tong, Yingqi Fan, Xuan Lu, Ping Nie, Wenjie Li, Xiaoyu Shen
Abstract
On-policy distillation (OPD) supervises student-generated trajectories with token-level teacher signals. Its sampled-token variant avoids the cost of full-vocabulary probabilities. Yet we find that training on fewer tokens can outperform full-token OPD, challenging the intuition that more supervision improves learning. This motivates selecting tokens by learning value. Existing disagreement-based criteria ignore probability scale: tokens assigned negligible probability by both models, termed low-low tokens, can receive large log-ratio rewards and hinder learning. We propose DIAL-OPD, a token-selection method that bridges log-probability and probability spaces by weighting reward magnitude with the logarithmic mean of teacher and student probabilities. A parameter beta controls this weighting, and the highest-scoring tokens are retained. Across 4 teacher-student pairs and 7 mathematical reasoning benchmarks, we compare DIAL-OPD with 9 baselines. Retaining only 40% of tokens, it outperforms Vanilla OPD and its full-token variants, with mean accuracy gains reaching 5.25 percentage points over Vanilla OPD, and doubles AIME25 Pass@16 from 13.33% to 26.67%. It also achieves up to an 18% relative improvement in mean accuracy over the strongest token-selection baseline at matched retention ratios. With a 4B teacher, DIAL-OPD surpasses the strongest full-token baseline using an 8B teacher at both student scales, showing that effective supervision allocation can outweigh teacher scaling. Further analysis shows that moderate beta balances suppressing low-low tokens against preserving useful disagreements. Token-level evidence reveals that DIAL-OPD filters high-reward tokens with limited reasoning value while preserving supervision critical to reasoning correctness.