Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall
Jacqueline He and colleagues at Meta AI, the University of Washington and Princeton show that standard knowledge distillation helps reasoning and hurts factual recall during mid-training, and trace the cause to teacher confidence.
Ask this paper
The stage dependence is the finding. Forward KL distillation improves both reasoning and factual recall during pre-training relative to next-token prediction, and during mid-training it continues to improve reasoning while slowing factual recall acquisition.
The cause is an asymmetry in teacher confidence. Teachers are more confident on procedural data than on knowledge-intensive data, and students acquire low-entropy factual knowledge earlier in training, so the distillation signal is strongest exactly where the student needs it least.
Switch Distillation uses teacher entropy as a routing signal. It distills on tokens where the teacher is confident and falls back to cross-entropy elsewhere, which is a one-line change to the objective.
It outperforms existing distillation objectives across teacher sizes, reaching 1.61 to 1.71x the reasoning performance and 1.13 to 1.19x the knowledge and commonsense performance relative to standard next-token prediction.
Directly actionable for anyone running a mid-training stage, because the default objective is silently trading factual recall for reasoning.
Abstract
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.