通过自适应调整蒸馏目标,提升语言模型生成多样性与对齐效果。
Entropy-Aware On-Policy Distillation of Language Models
- 引入前向KL补偿高熵教师分布下的信息损失,平衡精准模仿与覆盖多样性。
- 在六项数学推理任务上,小模型最高提升5.05%通过率,保持生成熵稳定。
- 适合追求生成质量与多样性的大模型蒸馏场景,尤其教师输出不确定时。
在策略蒸馏中,学生模型通过自身轨迹上的密集分词信号学习教师模型知识。传统方法使用反向KL散度,促使学生匹配教师的高置信预测,但其模式聚焦特性会降低生成多样性,并在教师分布熵较高时导致学习信号不稳定。为此,本文提出熵感知的在策略蒸馏方法:当教师熵较高时,动态加入前向KL项以捕捉合理输出的全范围,同时保留其他区域的精确模仿能力。该方法在不牺牲在策略训练效率的前提下,兼顾模式聚焦精度与模式覆盖鲁棒性。实验表明,该方法能维持生成熵(稳定分词级熵),并提升学生-教师对齐度(高熵分词上前向KL更低)。在六项数学推理基准测试中,相比基线方法,Qwen3-0.6B-Base、Qwen3-1.7B-Base和Qwen3-4B-Base分别获得+1.37、+2.39和+5.05的Pass@8准确率提升,证明考虑教师不确定性对维持多样性与实现有效知识迁移至关重要。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。