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
cs.CL
2026-09-02
On OLMo-2, forward-KL distillation lifts reasoning and facts together in pre-training but slows factual recall in mid-training. Switch Distillation distills only low-entropy teacher tokens, reaching 1.61–1.71× NTP reasoning while keeping 96.7–96.8% of factual recall.
Mid-training sits between pre-training and post-training: the same next-token objective, on a smaller, curated mix meant to lift reasoning, coding, and instruction following. Tokens are an order of magnitude fewer than pre-training, so each one is expensive, and distillation looks like a natural fit. A stronger teacher's predictive distribution is a richer target than a hard next-token label. Almost all distillation results were measured in pre-training or post-training. Whether the same recipe transfers to mid-training had not been checked.
On OLMo-2 it does not transfer. Forward-KL distillation improves reasoning and factual recall together during pre-training. During mid-training, reasoning keeps rising and factual acquisition slows. The paper calls this the reasoning–recall tradeoff.
The student is OLMo-2 1B. Pre-training starts from random weights for 100B tokens, past the Chinchilla optimum. Mid-training starts from the Stage-1 checkpoint already trained on 4T tokens and runs another 60B on Dolmino Mix 1124. Teachers are OLMo-2 Instruct at 1B, 7B, and 13B. Evaluation uses OLMES, grouped into reasoning, factual recall, and knowledge and commonsense macro-averages.
The mechanism work tracks three facts. Teacher predictive entropy is lower on procedural domains (math, instruction following) than on web knowledge, and within each domain lower entropy tracks higher top-1 agreement with the corpus token. Stratifying factual items by teacher entropy, the student has learned 67% of the lowest-entropy quintile by the end of pre-training and 5% of the highest. Unresolved facts at mid-training start are piled in the high-entropy bins, exactly where the teacher is least sure. Both forward and reverse KL shrink the gold-token logit gradient to about 0.5× NTP on those high-entropy facts.
Switch Distillation routes on teacher entropy: the lowest-entropy q=20% of in-batch tokens take reverse KL, the rest take cross-entropy. The routing signal is computed from teacher logits that distillation already needs, so there is no extra forward pass. Reverse KL is mode-seeking, which matches sharp, low-entropy teacher distributions.
After mid-training, Switch Distillation with a 7B teacher lifts reasoning from 26.1% under NTP to 44.7%; with a 13B teacher, to 42.1%. That is 1.71× and 1.61×. Factual recall is 29.3% against NTP's 30.3%, i.e. 96.7–96.8% retained. Knowledge and commonsense land at 49.3% and 46.5%, about 1.13–1.19×. Forward KL, reverse KL, and TRKD at α=0.5 all lean further into reasoning and further off recall. No KD operating point Pareto-dominates NTP during mid-training. The 7B teacher beats the 13B teacher overall, consistent with a capacity gap hurting distillation.
After the stock post-training stack (SFT, DPO, two RLVR rounds), reasoning reaches 50.6% with the 7B teacher, 1.32× NTP's 38.5% post-trained, and 48.0% with the 13B teacher, about 1.25×. The factual gap closes; Switch Distillation finishes slightly above NTP on recall, and knowledge and commonsense stay 1.13–1.20×. Ablations: swapping in forward KL costs 2.9 reasoning points; routing on teacher correctness, on domain, or at random all do worse; replacing soft distillation with teacher top-1 labels gains a bit of recall and drops reasoning 6.4 points. Entropy routing is the main driver, and the soft distribution is worth more than a hard label.
The same stage-dependent pattern shows up on SmolLM2. It is not an OLMo-2 recipe accident.
Mid-training is becoming a fixed stage in frontier recipes (Llama 3, SmolLM, OLMo-2). Pouring in pre-training-style distillation will systematically under-teach high-entropy facts. Switch Distillation is a near-zero-cost change: the same teacher logits, split by entropy into "trust the teacher" and "trust the corpus." For a team already running mid-training, this is an objective they can try without retuning the mix.
It asks how to extract supervision from a fixed pool, and sits next to work that edits the pool itself. The capacity gap is still real: a 7B teacher for a 1B student beats a 13B teacher.
Controlled runs are almost all OLMo-2 1B students, because almost no family releases intermediate checkpoints, the mid-training mix, and a full recipe together. SmolLM2 is a qualitative check. Teachers are Instruct models and the student is 1B; those choices were not grid-searched. Routing is a hard in-batch quantile switch, with q=20% from a small sweep, not a learned router. The authors guess late-stage pre-training may show the same tradeoff; this paper does not measure it. The factual group is TriviaQA, Natural Questions, and SimpleQA, which do not cover every shape of long-tail knowledge.