DIAL-OPD: Learning More from Fewer Tokens in On-Policy Distillation

2026-10-08

AI Summary

This study addresses the issue that blindly increasing supervised tokens in online policy distillation introduces noise, thereby hindering model learning. To overcome this, we propose a dynamic token selection algorithm based on learning value. Specifically, we design a scoring mechanism bridging logit and probability spaces, which employs reward-magnitude weighting to suppress low-probability noise and retain only high-value supervisory signals for efficient allocation. Our findings reveal that training with fewer tokens outperforms full supervision. Retaining merely 40% of tokens surpasses the baseline, improving average accuracy by 5.25% and doubling AIME scores. Notably, a smaller teacher model utilizing our selective approach outperforms a larger teacher model trained with all tokens.

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.

Problem

• On-policy distillation
• Token selection
• Knowledge distillation
• Low-probability tokens
• Mathematical reasoning

Innovation

• On-Policy Distillation
• Token Selection
• Mathematical Reasoning
• Knowledge Distillation
• Log-Probability Weighting