Distillation is the standard recipe for training small models: use a larger teacher model's output distribution as the supervision signal, and the student learns faster and better. But a paper submitted to arXiv on September 1 adds an important footnote to this common wisdom: distillation's effect depends on which training stage you apply it in — the same formulation behaves in exactly opposite ways during pre-training versus mid-training.

The phenomenon: one recipe, two stages, opposite results

The paper (arXiv:2609.01532) focuses on forward KL distillation — the standard form of knowledge distillation — with post-trained teachers. Through controlled experiments, the authors found:

  • During pre-training, forward KD improves both reasoning and factual recall simultaneously, relative to standard next-token prediction (NTP);
  • During mid-training — the intermediate phase of continued self-supervised learning on curated corpora — reasoning keeps improving, but factual recall acquisition actually slows down.

In other words, distillation is not a universal cure; it has stage-dependent side effects.

The cause: teacher confidence misaligned with the student's knowledge state

The authors trace this behavior to a combination of two factors. First, an asymmetry in teacher confidence across data domains: teachers are more confident on procedural data than on knowledge-intensive data. Second, the student's knowledge state evolves during training: students acquire low-entropy factual knowledge earlier in training. These two forces stack up, so that during mid-training the distillation signal actively interferes with factual knowledge.

Switch Distillation: route by confidence

The fix is surprisingly simple. The paper proposes Switch Distillation: distill only on tokens where the teacher is confident, using teacher predictive entropy as a lightweight routing signal; on all other tokens, fall back to cross-entropy. It effectively installs a switch on distillation — where the teacher is unsure, it simply doesn't teach.

Reported results:

  • Relative to standard NTP, it achieves 1.61-1.71x the reasoning performance and 1.13-1.19x on knowledge and commonsense, while preserving 96.7-96.8% of factual recall;
  • Across different teacher sizes, Switch Distillation consistently outperforms existing distillation objectives;
  • The benefits persist after post-training: reasoning retains 1.25-1.32x and knowledge/commonsense 1.13-1.20x gains, with the factual recall gap closed.

The paper spans 33 pages with 13 figures and 9 tables, with 12 authors including Luke Zettlemoyer and Wen-tau Yih; the code is open-sourced under the facebookresearch organization, in the midtraining-distillation repository.

So what

The value of this paper is not just another loss function — it decomposes "distillation" from a blanket term into a stage-aware operation: as more teams adopt mid-training phases today, indiscriminate distillation may be quietly damaging models' factual capabilities. If your training pipeline already includes this phase, routing by teacher entropy is a near-zero-cost improvement to validate. Full paper at https://arxiv.org/abs/2609.01532 .