中期训练阶段的知识蒸馏更偏向于推理而非事实回忆

摘要

基于Logit的知识蒸馏（KD）用于通过更强的教师模型监督来训练较小的语言模型（LM），但其收益在训练各阶段是否一致仍不清楚。通过受控实验，我们发现，使用后训练（post-trained）教师模型的前向Kullback-Leibler（KL）蒸馏——即标准KD形式——在中期训练（mid-training）阶段表现截然不同，而中期训练是自监督学习中介于预训练与后训练之间的一个中间阶段，在精选语料库上进行。令人惊讶的是，尽管前向KD在预训练阶段相较于标准下一词预测（NTP）能同时提升推理能力和事实回忆能力，但在中期训练阶段，它却在持续提升推理能力的同时减慢了事实回忆的获取。我们将这种阶段依赖性归因于教师模型在不同数据领域上的置信度不对称性以及学生模型不断演进的知识状态：教师模型在程序性数据上比在知识密集型数据上更为自信，而学生模型在训练早期便获取了低熵的事实性知识。为缓解这种不平衡，我们提出Switch Distillation——一种简单的中期训练目标，它在教师模型具有较高置信度的词元上进行蒸馏，使用教师预测熵作为轻量级路由信号，否则回退到交叉熵。Switch Distillation在不同教师规模下均一致优于现有蒸馏目标。相较于标准NTP，它实现了1.61–1.71倍的推理性能和1.13–1.19倍的知识与常识性能，同时保留了96.7–96.8%的事实回忆能力。关键在于，这些收益在后训练后仍然持续：Switch Distillation弥合了事实回忆差距，同时分别在推理以及知识与常识方面保持1.25–1.32倍和1.13–1.20倍的增益。

English

Logit-based knowledge distillation (KD) is used to train smaller language models (LMs) via supervision from stronger teachers, but whether its benefits are consistent across training stages remains unclear. Through controlled experiments, we find that forward Kullback-Leibler (KL) distillation--the standard KD formulation--with post-trained teachers behaves fundamentally differently during mid-training, an intermediate phase of self-supervised learning on curated corpora. Surprisingly, while forward KD simultaneously improves reasoning and factual recall during pre-training relative to standard next-token prediction (NTP), it instead slows factual recall acquisition during mid-training despite continued reasoning gains. We trace this stage dependence to an asymmetry in teacher confidence across data domains and the student's evolving knowledge state: teachers are more confident on procedural than knowledge-intensive data, while students acquire low-entropy factual knowledge earlier in training. To mitigate this imbalance, we propose Switch Distillation, a simple mid-training objective that distills on tokens where the teacher is confident, using teacher predictive entropy as a lightweight routing signal, and otherwise falls back to cross-entropy. Switch Distillation consistently outperforms existing distillation objectives across teacher sizes. Relative to standard NTP, it achieves 1.61-1.71x the reasoning performance and 1.13-1.19x the knowledge and commonsense performance while preserving 96.7-96.8% of factual recall. Crucially, these benefits persist after post-training: Switch Distillation closes the factual recall gap while maintaining 1.25-1.32x and 1.13-1.20x gains in reasoning and knowledge and commonsense, respectively.