知識蒸餾在中期訓練中更有利於推理而非事實記憶

摘要

基於logit的知識蒸餾（KD）常用於透過較強教師模型的監督來訓練較小的語言模型（LM），但其效益在不同訓練階段是否一致仍不清楚。透過受控實驗，我們發現採用後訓練教師的前向Kullback-Leibler（KL）蒸餾——即標準的KD形式——在中期訓練期間表現出根本不同的行為；中期訓練是指在精選語料上進行自監督學習的中間階段。令人意外的是，雖然前向KD在預訓練期間相對於標準下一個詞元預測（NTP）同時提升了推理與事實記憶能力，但在中期訓練期間，它卻在推理能力持續增進的同時減緩了事實記憶的習得。我們將這種階段依賴性追溯至教師信心在不同資料領域間的不對稱性，以及學生不斷演化的知識狀態：教師在程序性資料上的信心高於知識密集型資料，而學生則在訓練較早階段便獲取了低熵的事實知識。為緩解此失衡，我們提出Switch Distillation，這是一種簡單的中期訓練目標：它使用教師的預測熵作為輕量級路由信號，在教師信心較高的詞元上進行蒸餾，否則回退至交叉熵。無論教師規模大小，Switch Distillation都一致地優於現有的蒸餾目標。相對於標準NTP，它在保留96.7–96.8%事實記憶的同時，達到了1.61–1.71倍的推理表現，以及1.13–1.19倍的知識與常識表現。關鍵的是，這些效益在後訓練之後依然存在：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.