On-policy self-distillation improves math reasoning by reducing entropy overshoot
E$^2$-OPSD: Taming Entropy Overshoot in On-Policy Self-Distillation
Machine LearningArtificial Intelligence
Summary
Some AI models learn by teaching themselves, but this can cause confusion where the model’s guesses become too uncertain. The authors found that during training, a model’s confidence can become too low compared to its teacher, making learning less effective. They fixed this by having the model learn from similar solved problems and by adjusting how much it trusts itself based on confidence levels. This approach helps AI become better at math problems without needing extra computing resources.
What this means in practice
- •For ai developers: Improve AI models' ability to reason through math problems more accurately by reducing overconfidence and uncertainty during self-learning.
- •For language model trainers: Train language models with self-distillation techniques that better guide student models using similar examples rather than exact answers to enhance generalization.
Authors
Yifei Liu, Minghao Fang, Xinyu Gu, Chengkai Yao, Mengdi Liu, Tengfei Ma, Jiangbin Zheng, Chang Yu, Zhangyang Gao
Abstract
On-policy self-distillation (OPSD) provides dense token-level supervision without a second model: one network acts as teacher with the reference solution and as student with only the problem. We identify a specific failure mode of this recipe. During training, student token entropy rises past the teacher's and remains elevated, a pattern we call entropy overshoot. We trace it to both sides of distillation. The reference-conditioned teacher is confident along its answer-directed reasoning path, but this confidence transfers poorly to student-generated prefixes, making its supervision overly tied to answer-specific cues rather than reusable reasoning patterns; meanwhile, the forward KL used by OPSD continually diffuses the student's predictive distribution without pulling it back. We introduce E$^2$-OPSD to address both causes. Exemplar-guided teaching replaces the current answer with a retrieved solved neighboring problem, providing transferable reasoning guidance without revealing the destination and better matching student-reachable states. Entropy-aware distillation uses the student-teacher entropy gap to determine the direction and strength of each token's correction. E$^2$-OPSD improves math reasoning by up to 4.3 points in mean@16 over OPSD, while out-of-domain evaluations show gains over the corresponding base models of up to 4.9 points in mean@16 and 5.5 points in pass@8. Despite these gains, E$^2$-OPSD remains simple, requiring no additional forward passes or networks.