AL Alfred 2026-09-02 4 min read

Mid-Training Distillation: Reasoning 1.61x, Recall 96.7%

Knowledge distillation has a stage problem. The same forward-KL recipe that lifts both reasoning and factual recall during pre-training quietly erodes factual recall during mid-training — even as reasoning keeps climbing. That is the finding of a new Meta FAIR paper by Jacqueline He, Howard Yen, Shaurya Stella Li, Margaret Li, Hanqing Zeng, Yinglong Xia, Benyu Zhang, Zhuokao Zhao, Qiang Zhang, Pang Wei Koh, Luke Zettlemoyer, and Wen-tau Yih (arXiv:2609.01532, submitted 2026-09-01). The authors don't just document the failure mode — they ship a fix. Switch Distillation, a confidence-gated KL objective, reaches 1.61–1.71x the reasoning performance of plain next-token prediction while preserving 96.7–96.8% of factual recall, and the gains survive post-training.

Why Knowledge Distillation Fails During Mid-Training

Mid-training is the phase between pre-training and post-training where a model does self-supervised learning on curated, higher-quality corpora — the "continued pretraining" step most production stacks run. The paper's controlled experiments show that forward KL distillation (the standard KD objective, matching a student's distribution to a post-trained teacher) behaves fundamentally differently in this phase than in pre-training. During pre-training it improves both reasoning and factual recall relative to next-token prediction. During mid-training, recall growth stalls while reasoning gains continue.

The mechanism is an asymmetry in where the teacher's confidence lands. Post-trained teachers are markedly more confident on procedural data (math, code, stepwise reasoning) than on knowledge-intensive data (entities, facts). Meanwhile, students acquire low-entropy factual knowledge early in training. So by mid-training, the KL pressure is systematically strongest on reasoning tokens — and on the factual tokens, the teacher's weaker signal actively slows a student that already knows the answer. Blindly distilling everywhere, the paper shows, redirects the student's capacity toward procedure at the expense of memory.

Switch Distillation: Distill Only Where the Teacher Is Confident

The proposed fix is deliberately simple. Switch Distillation uses the teacher's predictive entropy as a per-token routing signal: distill (KL) on tokens where the teacher is confident, and fall back to standard cross-entropy everywhere else. No new architecture, no auxiliary heads, no reward model — just a threshold on teacher confidence at every token.

flowchart LR
    subgraph Token["Per-token routing"]
        T[Token] --> E{Teacher entropy low?}
        E -->|Yes| K[KL distil to teacher]
        E -->|No| C[Cross-entropy fallback]
    end
    K --> S[Student update]
    C --> S
    S --> R[Reasoning up 1.61x
Recall 96.7% kept]

That single mechanism outperforms existing distillation objectives consistently across teacher sizes, which is the paper's strongest empirical claim: the problem was never distillation itself, but where it was applied.

Results: Reasoning 1.61–1.71x, Factual Recall Held at 96.7%

Relative to standard next-token prediction during mid-training, Switch Distillation delivers 1.61–1.71x reasoning performance and 1.13–1.19x knowledge and commonsense performance, while preserving 96.7–96.8% of factual recall — versus plain forward KD, which sacrifices recall for the same reasoning gains. The benefits persist after post-training: Switch Distillation closes the factual recall gap while maintaining 1.25–1.32x reasoning and 1.13–1.20x knowledge and commonsense gains over the standard recipe. That the recall deficit is recoverable, not baked in, is the headline: the mid-training damage was a training-dynamics artifact, not a capability ceiling.

Limitations

Three caveats before anyone rewires a training pipeline. One: the stage-dependence cuts both ways — during pre-training, plain forward KD was strictly better, so the switch is a mid-training tool, not a universal one. Two: the routing signal is teacher entropy, a heuristic: a confidently-wrong teacher (a known failure mode at scale) would route more mass toward its errors, and the paper doesn't address teacher-correctness gating. Three: the evaluation is a controlled training study on curated corpora — 33 pages, 13 figures, 9 tables — but real production stacks bolt mid-training onto heterogeneous data mixes, and the exact crossover points (how much confidence, which entropy threshold) will need re-tuning per corpus.

Why Builders Should Care

Three practical takeaways. One: if you run any teacher-based continued pretraining, audit factual recall, not just reasoning benchmarks — this paper predicts exactly the regime where your KD loss is silently trading memory for procedure. Two: per-token confidence routing is nearly free (it's teacher entropy at inference time) and beats existing objectives across teacher sizes; it's the cheapest distillation upgrade available. Three: it reinforces the lesson from last week's OPD result: teacher supervision quality is the recurring failure point of distillation — there it was noise (50.6% wrong signals at scale), here it's confidence asymmetry. Measure where your teacher is confident before you let it teach.

Related reading