Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall

AI summary

Overview

• Research area: Natural Language Processing — specifically language model training methodology, with a focus on knowledge distillation (KD) during the mid-training stage of the training pipeline.
• Technical level: Intermediate. The paper uses standard distillation formalism (forward KL, reverse KL, cross-entropy, temperature scaling), but the conceptual framing — comparing training stages and diagnosing a capability tradeoff — is accessible to readers with basic familiarity with language model training.
• Scope: A controlled empirical study showing that logit-based knowledge distillation helps reasoning but hurts factual recall specifically during mid-training, an explanation of why, and a proposed routing objective called Switch Distillation that mitigates the tradeoff.

Affiliations listed: Meta AI, University of Washington, and Princeton University. Code is released at https://github.com/facebookresearch/midtraining-distillation. License: CC BY 4.0.

What This Paper Is About

Logit-based knowledge distillation trains a smaller student model to match the output distribution of a stronger teacher, and it has been studied mostly during pre-training and post-training. This paper asks whether those benefits carry over to mid-training — the intermediate phase where self-supervised next-token prediction continues on a smaller, curated, high-quality corpus — and finds that they do not behave the same way. The goal is to characterize that difference and design a training objective that recovers distillation's reasoning gains without sacrificing factual recall.

Key Contributions

1. Empirical finding: The authors uncover a robust "reasoning–recall tradeoff" during mid-training: distilling a substantially pre-trained student improves reasoning while slowing factual recall acquisition relative to next-token prediction (NTP). No KD objective tested Pareto-dominates NTP during mid-training.
2. Explanatory analysis: They trace the tradeoff to an interaction between teacher confidence, the student's evolving knowledge state, and the distillation objective — showing that facts not yet learned by the pre-trained student disproportionately receive weak teacher supervision.
3. A mid-training objective: They propose Switch Distillation, which routes each token between reverse-KL distillation and cross-entropy using teacher predictive entropy as a lightweight signal, with no additional parameters or model forward passes.
4. Persistence through post-training: They show the gains survive the full OLMo-2 post-training pipeline (SFT, DPO, and two rounds of RLVR), where Switch Distillation closes the factual recall gap while retaining reasoning and knowledge gains.

Main Findings

• Distillation strength shifts performance toward reasoning at both stages: Sweeping α ∈ [0.0, 0.3, 0.5, 0.7, 1.0], increasing the teacher's contribution generally moves the operating point toward higher reasoning performance across teacher sizes (1B, 7B, 13B Instruct), both KL directions, and both pre-training and mid-training, with diminishing returns at stronger distillation.
• The effect on factual recall is stage-dependent: During pre-training, moderate forward-KL distillation improves factual recall alongside reasoning, yielding Pareto improvements over NTP. During mid-training, no KD operating point outperforms NTP on factual recall even as reasoning improves. Because both stages use the same data mixture (Dolmino Mix 1124), the shift cannot be attributed to training data.
• Teacher supervision is asymmetric across domains: Procedural domains such as math and instruction-following concentrate at substantially lower teacher predictive entropy than unstructured text domains. Within every domain and teacher size, the probability that the teacher's top-1 prediction matches the ground-truth token decreases monotonically across entropy quintiles (Q1 lowest to Q5 highest).
• Low-entropy facts are learned first: By the end of pre-training, the student had learned 67% of facts in the lowest-entropy quintile (Q1) and only 5% in the highest (Q5). By mid-training initialization (after 4T tokens), unresolved facts are concentrated in the highest-entropy quintiles — exactly where teacher supervision is least confident.
• KD attenuates ground-truth supervision: Teacher probability on the ground-truth token strictly decreases with predictive entropy. Both forward and reverse KD increasingly attenuate the gold-token gradient as teacher entropy rises, reaching approximately 0.5× NTP for the highest-entropy facts (measured at α = 0.5). Both FKD and RKD incur their largest factual recall deficits on high-entropy facts.
• Switch Distillation wins during mid-training: With the OLMo-2 7B and 13B Instruct teachers, it raises the Reasoning macro-average from 26.1% under NTP to 44.7% and 42.1% respectively, and achieves the strongest Knowledge & Commonsense performance (49.3% and 46.5%), while staying closest to NTP on Factual Recall (29.3% and 29.3% versus 30.3%).
• Capacity gap matters: Distillation with the 7B teacher generally outperforms the 13B teacher, corroborating prior work that a large teacher–student size difference may reduce distillation effectiveness.
• Gains survive post-training: With the 7B teacher, Switch Distillation's Reasoning macro-average rises from 44.7% to 50.6% after post-training; with the 13B teacher, from 42.1% to 48.0%. It experiences the least forgetting and finishes with the highest Factual Recall macro-average, despite entering post-training with a small factual-recall deficit relative to NTP.
• Reported headline ratios versus NTP: Relative to NTP, Switch Distillation achieves 1.61–1.71× the reasoning performance and 1.13–1.19× the knowledge and commonsense performance while preserving 96.7–96.8% of factual recall during mid-training; after post-training it maintains 1.25–1.32× and 1.13–1.20× gains in reasoning and knowledge and commonsense respectively.
• Ablations isolate the mechanism: Replacing reverse KL with forward KL reduces Reasoning by 2.9 and Knowledge by 1.3 points; teacher-correct routing (-4.4/-5.1), random routing (-6.5/-0.8), and oracle domain routing (-7.2/-1.3) all underperform entropy-based routing; always using CE has little effect (-0.3/+0.1/-0.6); and teacher top-1 labels improve Factual Recall (+1.3) at the expense of Reasoning (-6.4) and Knowledge (-2.8).

Methodology in Plain English

The researchers built a controlled testbed on the fully open OLMo-2 ecosystem, which provides intermediate checkpoints, training recipes, and multiple model scales.

• Two training regimes on the same data. For pre-training, they initialize 1B-parameter students from random weights and train beyond Chinchilla-optimality for 100B tokens. For mid-training, they start from the OLMo-2 1B Stage 1 checkpoint (already pre-trained on 4T tokens) and continue for 60B tokens. Both use Dolmino Mix 1124, a mixture of filtered DCLM web text, FLAN instruction-following data, Dolmino Math, peS2o, Wikipedia (including Wikibooks), and Stack Exchange.
• Teachers. OLMo-2 1B Instruct, 7B Instruct, and 13B Instruct — post-trained teachers chosen because their instruction-following and reasoning abilities make them natural sources of capability transfer.
• Objectives compared. Standard next-token prediction; forward-KL distillation and reverse-KL distillation, each swept over the interpolation weight α; token-routing KD (TRKD), which applies forward-KL to high-entropy tokens while retaining CE everywhere; and the proposed Switch Distillation.
• Evaluation. The OLMES harness groups benchmarks into Reasoning (generative problem solving), Factual Recall (open-ended generative retrieval), Knowledge & Commonsense (multiple choice), and Instruction Following (post-training only). Macro-averages per group are reported. Factual recall is measured on TriviaQA, Natural Questions, and SimpleQA.
• Diagnostic analysis. The authors compute teacher predictive entropy for every token on sampled Dolmino documents, stratify examples into entropy quintiles, and track how much factual knowledge the student has acquired per quintile. They then derive the gold-token directional gradient under forward KD, reverse KD, and NTP to measure how strongly each objective reinforces the correct token.
• Switch Distillation's design. For each batch, the lowest q% of tokens by teacher entropy (q = 20%) are routed to reverse KL — chosen because its mode-seeking behavior suits sharp, low-entropy teacher distributions — while the remaining tokens fall back to cross-entropy. The two partitions are normalized separately so each contributes independently of its size. The routing signal comes from teacher logits already computed for distillation.
• Verification. Post-training applies OLMo-2 1B's four-stage pipeline (SFT, DPO, RLVR1, RLVR2) to each mid-trained model. Statistical significance over the strongest competing baseline is assessed with a paired bootstrap (p < 0.05). Supplementary checks use the SmolLM2 family (1.7B Instruct teacher, 360M student) and cross-family teacher analyses.

Why This Matters

The paper argues that distillation objectives should not be treated as stage-agnostic. Because mid-training is where models acquire much of the factual and reasoning prior that post-training later shapes, and because mid-training consumes orders of magnitude fewer tokens than pre-training while demanding far higher data quality, getting more learning signal out of each token matters. The finding that standard distillation silently trades away factual recall at this specific stage, and that a cheap entropy-based routing rule fixes it, gives practitioners a concrete, low-cost intervention rather than a data-side fix.

Real-world applications:

• Post-training pipelines: Teams that apply SFT, DPO, and RLVR to a mid-trained base can adopt the objective knowing it preserves reasoning gains and reduces forgetting of factual recall after alignment.
• Small-model deployment: Organizations training compact models (the paper's students are 1B parameters) for on-device or cost-sensitive inference can target better reasoning without regenerating or re-curating a training corpus.
• Data-constrained training: As high-quality human-written text becomes scarcer, extracting more supervision per observed token is a direct alternative to acquiring more data.
• Retrieval-augmented and knowledge-heavy systems: Models that must both reason and answer open-ended factual questions can benefit from a training recipe tuned to avoid the recall deficit.

Industry relevance: The work comes from Meta AI with collaborators at the University of Washington and Princeton University, uses the open OLMo-2 family and the OLMES evaluation harness, and is motivated explicitly by the use of post-trained teachers as reference models in frontier language modeling pipelines. The released code signal makes the method directly reproducible.

Future Directions

1. Stage-aware objectives beyond mid-training. The authors hypothesize that Switch Distillation may also help during late-stage or continual pre-training, once the student has already absorbed the easily transferred knowledge — but characterizing when standard KD stops yielding Pareto improvements is left open.
2. Smarter routing strategies. Teacher predictive entropy is a simple primitive. The paper suggests learned router networks, analogous to those used in Mixture-of-Experts architectures, might better capture when and where teacher supervision is beneficial.
3. Broader controlled validation. Isolating stage-dependent distillation effects requires intermediate checkpoints, matching data mixtures, and complete training recipes; few model families release all of these, and logit-based distillation additionally requires compatible vocabularies between student and teacher. The authors note their controlled experiments center on OLMo-2, with SmolLM2 and cross-family analyses as complementary tests.
4. Untested experimental dimensions. The authors do not exhaustively ablate choices such as the effect of teacher post-training or broader teacher scales, and they note that exploring these would require substantial additional compute.

Target Audience

This paper is most useful to researchers and engineers who train or fine-tune language models and make decisions about training objectives and supervision sources — particularly those working on mid-training, continued pre-training, data-efficient training, or teacher–student distillation. It also suits readers interested in the tradeoffs between reasoning and factual recall as distinct capabilities, and in how a model's existing knowledge state changes what supervision helps. Readers need basic familiarity with next-token prediction, cross-entropy, and Kullback–Leibler divergence to follow the derivations, though the core argument is presented through figures and controlled comparisons.

Authors’ abstract

Logit-based knowledge distillation (KD) is used to train smaller language models (LMs) via supervision from stronger teachers, but whether its benefits are consistent across training stages remains unclear. Through controlled experiments, we find that forward Kullback-Leibler (KL) distillation--the standard KD formulation--with post-trained teachers behaves fundamentally differently during mid-training, an intermediate phase of self-supervised learning on curated corpora. Surprisingly, while forward KD simultaneously improves reasoning and factual recall during pre-training relative to standard next-token prediction (NTP), it instead slows factual recall acquisition during mid-training despite continued reasoning gains. We trace this stage dependence to an asymmetry in teacher confidence across data domains and the student's evolving knowledge state: teachers are more confident on procedural than knowledge-intensive data, while students acquire low-entropy factual knowledge earlier in training. To mitigate this imbalance, we propose Switch Distillation, a simple mid-training objective that distills on tokens where the teacher is confident, using teacher predictive entropy as a lightweight routing signal, and otherwise falls back to cross-entropy. Switch Distillation consistently outperforms existing distillation objectives across teacher sizes. Relative to standard NTP, it achieves 1.61-1.71x the reasoning performance and 1.13-1.19x the knowledge and commonsense performance while preserving 96.7-96.8% of factual recall. Crucially, these benefits persist after post-training: Switch Distillation closes the factual recall gap while maintaining 1.25-1.32x and 1.13-1.20x gains in reasoning and knowledge and commonsense, respectively.