Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall

Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall. By Jacqueline He and eleven others, from Meta AI, the University of Washington, and one other institution. Published on the first of September. Its code repository is worked out: about forty-five source files, with a readme and a licence.

Logit-based knowledge distillation trains a smaller language model using supervision from a stronger teacher. Whether its benefits stay the same at every stage of training has not been clear. The claim here is that they do not. Standard forward Kullback-Leibler distillation, with post-trained teachers, behaves fundamentally differently during mid-training, the intermediate phase of self-supervised learning on curated corpora. During pre-training it improves both reasoning and factual recall compared with plain next-token prediction. During mid-training, reasoning keeps gaining, but the acquisition of factual recall slows down.

The evidence comes from controlled experiments, and the authors trace the stage dependence to an asymmetry with two sides. On the teacher's side, confidence differs by domain: teachers are more confident on procedural data than on knowledge-intensive data. On the student's side, knowledge evolves over training: students pick up low-entropy factual knowledge early. The paper reads the slowdown in factual recall as the meeting point of those two facts: an imbalance between what the teacher is sure of and what the student is learning at that moment.

To test that account, they propose Switch Distillation, a simple mid-training objective. It distills only on tokens where the teacher is confident, using the teacher's predictive entropy as a lightweight routing signal, and falls back to ordinary cross-entropy everywhere else. It consistently beats existing distillation objectives across teacher sizes.

Against standard next-token prediction, Switch Distillation reaches between one point six one and one point seven one times the reasoning performance. On knowledge and commonsense it reaches between one point one three and one point one nine times. And it keeps ninety-six point seven to ninety-six point eight percent of factual recall.

The benefits survive post-training. After it, Switch Distillation closes the factual recall gap, while holding gains of one point two five to one point three two times in reasoning, and one point one three to one point two times in knowledge and commonsense.

The finding is about a particular stage and a particular formulation. It concerns forward distillation with post-trained teachers, and the trouble is specific to mid-training: in pre-training the same method helps both abilities. The proposed fix addresses the imbalance through routing by teacher confidence, and its results are reported relative to standard next-token prediction and to other distillation objectives.

Distillation, on this account, is not uniformly good or bad. It depends on when you apply it, and on whether the teacher is actually sure of what it is passing on.