Entropy-Aware On-Policy Distillation of Language Models

On-policy distillation is a promising approach for transferring knowledge between language models, where a student learns from dense token-level signals along its own trajectories. This framework typically uses reverse KL divergence, encouraging the student to match the teacher's high-confidence predictions. However, we show that the mode-seeking property of reverse KL reduces generation diversity and yields unstable learning signals when the teacher distribution has high entropy. To address this, we introduce Entropy-Aware On-Policy Distillation. Our key idea is augmenting the standard reverse KL objective with forward KL when teacher entropy is high, capturing the full range of plausible outputs while retaining precise imitation elsewhere. It balances mode-seeking precision with mode-covering robustness without sacrificing on-policy training efficiency. Experiments show that our method maintains generation diversity (sustained token-level entropy) and improves student–teacher alignment (lower forward KL on high-entropy tokens). Across six math reasoning benchmarks, this yields Pass@8 accuracy gains of $+1.37$ for Qwen3-0.6B-Base, $+2.39$ for Qwen3-1.7B-Base, and $+5.05$ for Qwen3-4B-Base compared to baseline on-policy distillation methods. These results demonstrate that accounting for teacher uncertainty is essential for maintaining diversity and achieving effective knowledge transfer. Our code is publicly available at https://github.com/WLS04/EOPD.

Large AI language models are powerful but heavy and expensive to run, so a technique called "distillation" is widely used to transfer their abilities into smaller, more efficient models. Existing methods train the small model (the "student") to imitate only the most confident answers of the large model (the "teacher"). The problem appears in hard tasks like math, where many different solution paths are valid: even when the teacher deliberately keeps several possibilities open, the student fixates on just one. As a result, the student loses its ability to explore diverse solutions, and its training becomes unstable. We taught the model to tell apart the moments when the teacher is confident from the moments when it is hesitant. When the teacher is confident, the student learns to follow its answer directly, as before. But when the teacher is weighing several possibilities, the student learns to preserve that full range of options instead of collapsing onto one. Because this extra guidance is applied only when it is needed, our method keeps training efficient while transferring the teacher's richer reasoning far more faithfully. Small models trained this way explore diverse solutions better and consistently solve more math reasoning problems correctly. This helps build compact, efficient AI with strong reasoning ability without relying on large, costly models, which in turn can reduce the computational and environmental cost of using AI.