Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall
TL;DR - Standard forward-KL knowledge distillation during language-model mid-training improves reasoning but can slow factual recall acquisition. The proposed entropy-routed Switch Distillation preserves most factual recall while delivering stronger reasoning and knowledge performance than next-token prediction.
- Teacher confidence is higher on procedural data than knowledge-intensive data, creating stage-dependent distillation effects as student knowledge evolves.
- Switch Distillation uses teacher predictive entropy to distill confident tokens and applies cross-entropy elsewhere.
- It achieves 1.61–1.71× reasoning and 1.13–1.19× knowledge and commonsense performance while retaining 96.7–96.8% of factual recall versus standard next-token prediction.
- After post-training, it closes the factual-recall gap while retaining 1.25–1.32× reasoning and 1.13–1.20× knowledge and commonsense gains.