Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall

TL;DR

中期训练中的知识蒸馏提升推理能力但减缓事实回忆,提出Switch Distillation方法。

cs.CL 🔴 高级 2026-09-02 35 次浏览
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
知识蒸馏 推理 事实回忆 中期训练 语言模型

核心发现

方法论

研究采用前向KL蒸馏和Switch Distillation方法,分析教师模型在中期训练中的信心不对称性。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倍的知识提升。
  • 实验表明,Switch Distillation在不同教师规模下均优于现有蒸馏目标。

研究意义

该研究揭示了知识蒸馏在中期训练中的独特行为,提出的Switch Distillation方法有效缓解了推理与事实回忆之间的权衡。这一发现对语言模型的训练策略提供了新的视角,特别是在数据稀缺的情况下,提升了模型的推理能力。

技术贡献

技术贡献包括提出Switch Distillation方法,利用教师预测熵进行路由,显著改善了中期训练中的推理性能,同时保持了事实回忆的完整性。这一方法为中期训练阶段的知识蒸馏提供了新的思路。

新颖性

首次揭示了知识蒸馏在中期训练中对推理与事实回忆的不同影响,并提出了Switch Distillation作为解决方案,与传统方法相比具有显著优势。

局限性

  • Switch Distillation在高熵情况下可能无法充分利用教师模型的指导,导致事实回忆的轻微下降。
  • 该方法在特定领域数据上的表现可能受限于教师模型的预测熵。
  • 在极端数据稀缺的情况下,效果可能不如预期。

未来方向

未来研究可以探索如何在更广泛的数据集和任务上应用Switch Distillation,进一步优化其在不同训练阶段的性能。此外,研究如何结合其他蒸馏策略以增强模型的整体能力也是一个值得探索的方向。

AI 总览摘要

在语言模型的训练中,知识蒸馏是一种常用的方法,通过强大的教师模型来指导较小的学生模型。然而,这种方法在中期训练阶段的效果一直存在争议。本文通过实验发现,传统的前向KL蒸馏在中期训练中对推理能力有显著提升,但对事实回忆的获取速度有所减缓。

为了解决这一问题,研究者提出了Switch Distillation方法,该方法利用教师模型的预测熵来决定是否进行蒸馏。具体来说,当教师模型对某个数据点的预测信心较高时,进行蒸馏;否则,回退到交叉熵监督。这种方法在实验中表现出色,在推理和知识获取方面均优于传统方法。

Switch Distillation的提出不仅为中期训练阶段的知识蒸馏提供了新的解决方案,还为未来的研究指明了方向。尽管该方法在某些情况下可能面临挑战,但其在提升推理能力的同时保持事实回忆的完整性,显示出巨大的潜力。

深度分析

研究背景

近年来,语言模型的训练逐渐依赖于知识蒸馏技术,通过强大的教师模型来指导学生模型的学习。然而,知识蒸馏在不同训练阶段的效果尚未得到充分研究,尤其是在中期训练阶段。中期训练通常涉及在高质量数据集上的自监督学习,旨在提升模型的推理和事实性能力。

核心问题

核心问题在于传统的知识蒸馏方法在中期训练阶段对推理和事实回忆的影响不一致。具体来说,虽然推理能力有所提升,但事实回忆的获取速度却减缓。这种不一致性可能源于教师模型在不同数据域上的信心不对称性。

核心创新

本文的核心创新在于提出了Switch Distillation方法,通过教师模型的预测熵来路由蒸馏过程。与传统方法不同,该方法在教师信心高的情况下进行蒸馏,否则回退到交叉熵监督,从而有效缓解了推理与事实回忆之间的权衡。

方法详解

  • �� 使用前向KL蒸馏进行初步实验,分析其在中期训练中的表现。
  • �� 提出Switch Distillation方法,通过计算教师模型的预测熵来路由蒸馏过程。
  • �� 在教师信心高的情况下进行蒸馏,否则回退到交叉熵监督。
  • �� 在多个教师规模下进行实验,验证方法的有效性。

实验设计

实验使用OLMo-2生态系统,包含1B学生模型和7B、13B教师模型。数据集包括Dolmino Mix 1124等,实验设计涵盖了不同的蒸馏强度和KL方向。通过对比标准NTP和不同蒸馏方法,评估Switch Distillation的性能。

结果分析

实验结果显示,Switch Distillation在推理性能上比标准NTP提高1.61-1.71倍,知识和常识性能提高1.13-1.19倍,同时保留96.7-96.8%的事实回忆。这些提升在后期训练后仍然存在,显示出该方法的稳健性。

应用场景

Switch Distillation可用于需要高推理能力的语言模型训练场景,如智能助手和自动问答系统。其在数据稀缺情况下的优势使其适用于资源有限的应用场景。

局限与展望

尽管Switch Distillation在推理能力上表现出色,但在高熵情况下可能无法充分利用教师模型的指导,导致事实回忆的轻微下降。此外,该方法在特定领域数据上的表现可能受限于教师模型的预测熵。

通俗解读 非专业人士也能看懂

想象一下,你在厨房里做饭。传统的知识蒸馏就像是跟着一本食谱做菜,食谱告诉你每一步该怎么做。但在中期训练中,这个食谱有时候不太适合,因为它对某些步骤的信心不够高。Switch Distillation就像是一个聪明的助手,它会根据每一步的难度和重要性来决定是否跟随食谱。如果某一步很简单,它会让你自己做;如果很复杂,它会给你更多的指导。这样,你不仅能做出美味的菜肴,还能学到更多的烹饪技巧。

简单解释 像给14岁少年讲一样

嘿,小伙伴们!想象一下你在玩一个超级复杂的游戏,游戏里有个超级厉害的NPC(非玩家角色)教你怎么打怪升级。通常情况下,这个NPC会告诉你每一步该怎么做,但有时候它也会犯错,尤其是在一些特别难的关卡。Switch Distillation就像是一个聪明的游戏助手,它会根据每个关卡的难度来决定是否让NPC指导你。如果关卡很简单,它会让你自己探索;如果很难,它会给你更多的提示。这样,你不仅能通关,还能成为游戏高手!

术语表

知识蒸馏 (Knowledge Distillation)

一种通过强大的教师模型指导较小学生模型学习的方法,通常用于提升模型性能。

在本文中用于提升语言模型的推理能力。

前向KL蒸馏 (Forward KL Distillation)

一种通过最小化教师和学生预测分布之间的前向KL散度来进行蒸馏的方法。

在实验中用于分析中期训练的效果。

Switch Distillation

一种利用教师预测熵来路由蒸馏过程的新方法,旨在提升中期训练的推理性能。

本文提出的核心方法,用于解决推理与事实回忆的权衡问题。

预测熵 (Predictive Entropy)

衡量模型对某个预测的信心程度,熵越低表示信心越高。

用于决定是否进行蒸馏。

Dolmino Mix 1124

一个包含多种数据源的数据集,用于训练和评估语言模型。

在实验中用于测试Switch Distillation的性能。

开放问题 这项研究留下的未解疑问

  • 1 如何在更多样化的数据集上应用Switch Distillation以验证其通用性。
  • 2 探索结合其他蒸馏策略以进一步提升模型性能的可能性。
  • 3 研究在极端数据稀缺情况下Switch Distillation的适用性。

应用场景

近期应用

智能助手

通过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.

cs.CL