Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall

TL;DR Switch Distillation 在 mid-training 阶段按教师预测熵选择性路由，仅在教师自信处做 logit 蒸馏，其余回退交叉熵，从而在提升推理的同时保留事实记忆，解决标准 KL 蒸馏在 mid-training 对事实获取的抑制。

问题

问题背景

当前小型语言模型训练普遍采用 logit-based knowledge distillation (KD),由教师模型提供 logits 监督。mid-training 作为预训练后的中间阶段,用于在 curated corpora 上增强领域知识与推理能力。

现有方法局限

标准 forward KL distillation 在所有训练阶段统一应用,但论文发现其效果存在阶段依赖性:pre-training 阶段,forward KD 同时提升推理和事实 recall;mid-training 阶段,却显著拖慢事实获取。原因包括:教师置信度在程序性数据上偏高、在知识密集型数据上偏低;且学生此时已获得低熵事实知识,KD 会削弱对事实 token 的监督。因此固定 KD 目标无法平衡推理与事实召回,工程上直接沿用 pre-training 设置会导致模型事实能力倒退。

为什么这个问题难/重要

mid-training 数据混合事实与推理,若蒸馏策略不当,会牺牲事实准确,影响下游 QA 等任务。业界追求小型模型同时具备强推理与可靠事实记忆,单纯 NTP 或全局 KD 难以兼顾。论文提出 teacher predictive entropy 作为路由信号,但需要在每个 token 级做自适应决策,训练开销与稳定性是挑战。与 reverse KL、数据筛选等方法相比,其路由开销更低,但动态切换目标的实现复杂度更高。

行业类比

类似设备端语音助手蒸馏,若 mid-training 阶段 KD 策略错误,模型“会推理但记不住事实”,用户询问具体知识时容易出错。

核心洞察

• 知识蒸馏的效果高度依赖训练阶段：在 pre-training 中 forward KL 同时提升推理与事实召回，但在 mid-training 中却会损害事实召回，尽管推理仍持续增益。这一发现打破了“蒸馏收益单调一致”的假设，指出学生模型已有的知识状态会改变蒸馏信号的相对价值——当学生已掌握低熵事实知识时，教师在此类 token 上的监督反而成为干扰，导致事实遗忘。与以往仅关注蒸馏强度或教师规模的工作不同，该研究首次系统刻画了阶段依赖的 tradeoff，并为分阶段设计蒸馏目标提供了直接依据。
• Switch Distillation 使用教师预测熵作为轻量级路由信号，仅在教师高置信度 token 上蒸馏，其余位置回退到标准交叉熵，从而在 mid-training 中同时收获推理增益并保留事实召回。这一设计的独特性在于：它不依赖额外数据标注或复杂课程策略，而是利用教师模型自身的置信度分布自动识别适合蒸馏的 token 类型（高熵的程序性/推理数据 vs 低熵的知识密集型数据），有效缓解了 forward KD 在 mid-training 阶段对事实监督的稀释效应。相比现有方法固定全量蒸馏或单纯加权重，该路由机制能以极低开销（计算教师熵）实现 1.6x 推理性能与近乎无损的事实召回。

方法

Switch Distillation 是一种面向 mid-training 阶段的选择性知识蒸馏方法，用于在蒸馏强教师时缓解事实性知识召回下降的问题。

输入与背景

• 学生模型：较小的语言模型，处于 mid-training 阶段（在精选语料上进行自监督学习）。
• 教师模型：经过 post-training 的更大模型，提供 logits 监督。
• 标准做法：forward KL 蒸馏在所有 token 上统一使用教师分布作为目标。

关键模块

1. 教师预测熵计算：对每个 token，计算教师输出分布的熵，作为置信度的轻量代理指标。
2. Token 级路由：设定一个熵阈值（或分位数 q），当教师熵低于阈值时认为教师“有把握”，在该 token 上使用 KL 蒸馏损失；否则回退到标准 交叉熵（NTP）。
3. 损失切换：损失函数在蒸馏和交叉熵之间按 token 动态切换，无需显式数据域标签。

输出与效果

• 学生模型在推理能力上获得与标准蒸馏相当的提升，同时事实回忆得以保留（相对 NTP 保留约 96.7-96.8%）。
• 在 post-training 后，事实回忆差距进一步缩小，推理和知识/常识指标仍保持显著增益。

核心洞察：教师在过程性数据上置信度高、在知识密集型数据上置信度低，而学生早期已获得低熵事实知识；统一蒸馏会削弱事实监督，选择性蒸馏则避免这一问题。

与同类方法的差异

与在整个训练阶段对所有 token 施加 forward KL 的标准蒸馏不同，Switch Distillation 利用教师预测熵实现动态路由，只在教师高置信度 token 上蒸馏，从而在提升推理的同时保留事实知识获取。

实验

实验设计

论文在两种训练阶段( pre-training 和 mid-training )下对比了标准 forward KL 蒸馏(即 logit-based KD)与标准 next-token prediction (NTP) 的效果,教师模型为 post-trained teachers。评估涵盖推理、知识/常识和事实回忆三个维度。基于观察到的训练阶段依赖性,提出 Switch Distillation,其利用 teacher predictive entropy 作为轻量路由信号:仅对教师高置信度的 token 进行蒸馏,否则回退到 cross-entropy。

关键发现

• 在 pre-training 阶段,forward KD 同时提升推理和事实回忆;但在 mid-training 阶段,它虽然继续提升推理,却减慢事实回忆获取。
• 根因是教师在不同数据域上的置信度不对称(对程序性数据比知识密集型数据更自信),而学生在训练早期已获得低熵事实知识,导致 KD 削弱了事实监督。
• Switch Distillation 有效缓解此不平衡:与标准 NTP 相比,mid-training 阶段推理性能提升 1.61-1.71x,知识/常识提升 1.13-1.19x,事实回忆保留 96.7-96.8%。
• 经 post-training 后,事实回忆差距闭合,推理保持 1.25-1.32x 提升,知识/常识保持 1.13-1.20x 提升。

与基线对比解读

标准 forward KD 在 mid-training 阶段对事实回忆有害,而 Switch Distillation 通过教师置信度路由避免了这一缺陷,且计算开销极低(仅需教师输出熵)。与 reverse KL 等目标相比,论文指出二者优化几何不同,但 Switch 方法更简单且无需额外模型。对工程实践的启示是:在 mid-training 阶段应用蒸馏时,应根据学生知识状态和教师置信度采用自适应路由,而非全局均匀蒸馏。

行业影响

落地场景

Switch Distillation 在 mid-training 阶段蒸馏小模型，提升推理能力同时保持事实召回。适合需要 低成本推理 的场景：电商智能客服中的多步推理问答、内容平台的摘要与事实核查、医疗辅助诊断中的指南遵循、金融报告生成等。结合 teacher predictive entropy 路由，可训练 1-3B 参数模型替代 7B+ 模型，部署到移动端或边缘设备。

商业价值

• 降本：小模型推理成本显著低于大模型；Switch Distillation 后的小模型在 reasoning 上达到 1.61-1.71x NTP 基线，factual recall 保留 96.7-96.8%，可减少对大规模 API 的依赖。
• 体验提升：保留事实准确性，避免“幻觉”导致用户不信任；推理能力增强可处理复杂指令，提升自动化率。
• 增收：通过提高客服/推荐等场景的自动化水平，释放人力资源。

与现有工作流集成

可嵌入标准 LLM 训练 pipeline：在 mid-training 阶段加入 Switch Distillation 损失，替换常规 NTP；现有 post-training（SFT/RLHF）不变。实现简单，只需计算 teacher predictive entropy 并用阈值路由。已在 Meta AI 的 midtraining-distillation 开源，可基于 HuggingFace Trainer 修改。对多模型蒸馏的企业，可结合现有 teacher 模型作为服务。

局限

• **实验规模与模型家族的局限性**：论文主要在较小规模的学生模型（约 1B 参数级别）和有限的 post-trained 教师模型上进行验证，数据集偏向特定语料，这限制了结论向更大规模模型（如 7B+）或不同架构（如 MoE）的外推性。工程上，将 Switch Distillation 应用于生产级大模型前，需自行验证教师置信度分布是否仍具有不对称性，以及路由信号是否有效。与 LoRA 等通用方法相比，该方法的验证范围较窄。
• **Switch Distillation 的超参依赖性**：方法引入教师熵阈值作为路由信号，但论文虽进行了部分消融（如 q 的取值），实际部署时仍需要对阈值进行调优，且最优阈值可能随数据分布和训练阶段变化。此外，教师置信度本身可能不是完美的监督信号，低置信度样本不一定没有蒸馏价值，强制 fallback 到 cross-entropy 可能损失部分有用但教师不确定的知识。工程上需要额外监控路由比例和训练动态。
• **仅针对 mid-training 阶段的解决方案**：论文发现 forward KD 在 mid-training 阶段损害 factual recall，并提出 Switch Distillation 缓解，但并未探索该方法在 pre-training 或 post-training 阶段是否同样有效，也未讨论与其他 KD 变体（如 reverse KL）的完整对比，虽然附录中分析了 forward 与 reverse KL 的几何差异，但主实验未将 Switch Distillation 与 reverse KL 直接对比。这限制了其作为通用训练策略的适用性。