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.
翻译:在线策略蒸馏是一种有前景的语言模型知识迁移方法,其中学生模型沿着自身轨迹接收密集的标记级信号进行学习。该框架通常采用反向KL散度,促使学生模型匹配教师模型的高置信度预测。然而,我们发现教师分布高熵时,反向KL的模式寻求特性会降低生成多样性并产生不稳定的学习信号。为解决此问题,我们提出熵感知在线策略蒸馏。核心思想是在教师熵较高时用前向KL增强标准反向KL目标,在保持精确模仿的同时捕获全部合理输出范围。该方法在不牺牲在线策略训练效率的前提下,平衡了模式寻求的精确性与模式覆盖的鲁棒性。实验表明,我们的方法维持了生成多样性(稳定的标记级熵),并改善了学生-教师对齐(高熵标记上前向KL降低)。在六个数学推理基准上,相比基线在线策略蒸馏方法,Pass@8准确率在Qwen3-0.6B-Base提升+1.37,Qwen3-1.7B-Base提升+2.39,Qwen3-4B-Base提升+5.05。这些结果证明,考虑教师不确定性对于维持多样性并实现有效知识迁移至关重要。