Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall
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
Why it is worth reading
This paper uncovers a critical trade-off in language model training: applying standard knowledge distillation during mid-training accelerates reasoning capabilities but actively impedes the acquisition of factual recall. The authors diagnose the mechanism behind this phenomenon—namely, that teachers exhibit high predictive uncertainty on factual tokens, thereby diluting the ground-truth training signal for facts that the student has not yet acquired. To resolve this, they introduce Switch Distillation, a simple and computationally lightweight objective that routes confident tokens to reverse-KL distillation and uncertain tokens to standard cross-entropy. For practitioners developing foundation models, this provides a practical, drop-in training objective that substantially improves reasoning performance without sacrificing factual knowledge, with benefits that persist through subsequent alignment and post-training pipelines.
Core research claims
- Mid-training ablation results for Always CE relative difference across Reasoning, Factual Recall, and Knowledge task groups.
- Downstream task evaluation after post-training stage RLVR1 for 13b_rkd.
- The stage-dependent distillation tradeoff generalizes to the SmolLM family (SmolLM2 360M student, 1.7B Instruct teacher, SmolLM3 Stage-3 data), where mid-training KD exhibits the reasoning-recall deficit and Switch Distillation mitigates it.