Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall

📄 arXiv: 2609.01532v1 📥 PDF

作者: Jacqueline He, Howard Yen, Shuyue Stella Li, Margaret Li, Hanqing Zeng, Yinglong Xia, Benyu Zhang, Zhuokai Zhao, Qiang Zhang, Pang Wei Koh, Luke Zettlemoyer, Wen-tau Yih

分类: cs.CL

发布日期: 2026-09-01

备注: 33 pages, 13 figures, 9 tables. Code is publicly available at https://github.com/facebookresearch/midtraining-distillation

💡 一句话要点

提出Switch Distillation以解决中期训练中的知识蒸馏问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 知识蒸馏 语言模型 推理能力 事实回忆 自监督学习 模型训练 教师模型 学生模型

📋 核心要点

  1. 现有的知识蒸馏方法在中期训练阶段表现不一致，导致推理与事实回忆之间的平衡失调。
  2. 论文提出Switch Distillation，通过教师的预测熵来选择自信的标记进行蒸馏，从而改善中期训练效果。
  3. 实验表明，Switch Distillation在推理和知识表现上均有显著提升，同时保持了高水平的事实回忆能力。

📝 摘要（中文）

基于logit的知识蒸馏（KD）用于通过强教师对较小语言模型（LMs）进行监督训练，但其在不同训练阶段的效果尚不明确。通过控制实验，我们发现标准的前向Kullback-Leibler（KL）蒸馏在中期训练阶段表现出根本不同的行为。尽管前向KD在预训练阶段相较于标准的下一个标记预测（NTP）同时提高了推理和事实回忆，但在中期训练阶段却减缓了事实回忆的获取。为了解决这一不平衡，我们提出了Switch Distillation，这是一种简单的中期训练目标，通过教师的预测熵作为轻量级路由信号，在教师自信的标记上进行蒸馏。实验结果表明，Switch Distillation在各类教师模型中均优于现有蒸馏目标。

🔬 方法详解

问题定义：论文要解决的问题是现有知识蒸馏方法在中期训练阶段的效果不一致，尤其是推理与事实回忆之间的平衡问题。现有方法在中期训练中可能导致事实回忆的减缓。

核心思路：论文的核心思路是提出Switch Distillation，通过教师的预测熵来选择自信的标记进行蒸馏，避免在低自信标记上进行蒸馏，从而改善中期训练的效果。

技术框架：整体架构包括教师模型和学生模型，教师模型提供预测，学生模型通过Switch Distillation进行训练。具体流程为：在教师自信的标记上进行蒸馏，其他情况下使用交叉熵损失进行训练。

关键创新：最重要的技术创新点是引入了教师预测熵作为轻量级路由信号，以选择自信的标记进行蒸馏。这一设计使得模型在中期训练阶段能够更有效地学习。

关键设计：关键设计包括损失函数的选择，使用了基于教师预测熵的蒸馏损失和交叉熵损失的结合。此外，模型的训练过程中，教师和学生的知识状态的动态调整也是一个重要的设计考虑。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果显示，Switch Distillation相较于标准的下一个标记预测（NTP），在推理性能上提升了1.61-1.71倍，知识和常识性能提升了1.13-1.19倍，同时保持了96.7-96.8%的事实回忆能力。这些优势在后期训练中依然持续存在。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和智能问答等。通过提高模型的推理能力和知识表现，Switch Distillation能够为实际应用提供更高效的语言理解和生成能力，具有重要的实际价值和未来影响。

📄 摘要（原文）

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.