PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2609.01532v1cs.CL

TL;DR

The paper asks whether knowledge distillation benefits language-model training consistently across pre-training and mid-training. Through controlled experiments, it explains the stage-dependent reasoning–recall tradeoff and introduces Switch Distillation, which routes tokens by teacher entropy. Switch Distillation improves reasoning and knowledge-related performance while largely preserving factual recall, with gains persisting after post-training.

  • Problem

    Evidence is limited on whether logit-based knowledge distillation benefits remain consistent across language-model training stages, especially during mid-training.

  • Method

    Switch Distillation uses teacher predictive entropy to route confident tokens to distillation and other tokens to cross-entropy during mid-training.

  • Results

    Switch Distillation consistently improves the reasoning–recall tradeoff across teacher sizes and closes the factual recall gap after post-training while retaining capability gains.

  • Takeaways & Limitations

    Distillation should be stage-aware and selective because teacher uncertainty and the student’s evolving knowledge state affect which tokens benefit from teacher supervision.

  • Takeaways & Limitations

    The study focuses on mid-training and suggests, rather than establishes, that its principle may extend to late-stage or continual pre-training.

Abstract

from arXiv · show

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.

1 Introduction

The paper finds that knowledge distillation is stage-dependent: it improves reasoning across training, but standard distillation slows factual acquisition during mid-training. The authors explain this tradeoff and propose Switch Distillation to route tokens using teacher confidence.

  • This stage dependence contrasts with pre-training, where forward-KL distillation improves reasoning and factual recall together.
  • The tradeoff reflects interactions among teacher confidence, the student’s evolving knowledge, and the distillation objective.
  • After post-training, Switch Distillation retains reasoning and knowledge gains while closing the factual-recall gap.With an OLMo-2 7B Instruct teacher, reasoning is 32% higher, knowledge tasks improve by 20%, and the factual recall gap closes entirely.
  • Mid-training distillation improves reasoning while slowing factual acquisition relative to next-token prediction.
  • Switch Distillation routes tokens between knowledge distillation and next-token prediction using teacher predictive entropy.It is presented as a mid-training objective designed to mitigate the reasoning–recall tradeoff.

2 Background

The paper frames mid-training as continued self-supervised learning on curated data and introduces logit-based distillation as an alternative to standard next-token prediction. Experiments use the OLMo-2 ecosystem and grouped benchmark evaluation.

  • Logit-based knowledge distillation trains a student to match a vocabulary-compatible teacher’s output distribution.
  • Standard KD uses forward-KL divergence, while the paper also studies reverse-KL distillation.
  • The experiments use OLMo-2 Instruct teachers at 1B, 7B, and 13B scales and train on the Dolmino Mix 1124 data mixture.
  • Evaluation groups tasks into Reasoning, Factual Recall, Knowledge & Commonsense, and post-training Instruction Following.The stage-dependent analysis focuses on Reasoning and Factual Recall.

3 Characterizing Mid-Training Distillation

Distillation consistently shifts performance toward reasoning, but its effect on factual recall depends on the training stage. During mid-training, no KD operating point outperforms NTP on factual recall, whereas post-training preserves reasoning gains and can restore recall.

  • Increasing distillation strength generally raises reasoning performance across teacher sizes, KL directions, and training stages, with diminishing returns.
  • During pre-training, moderate forward-KL distillation improves factual recall alongside reasoning and can Pareto-improve over NTP.
  • During mid-training, no KD operating point outperforms NTP on factual recall even as reasoning improves.
  • After identical post-training, factual-recall deficits narrow or reverse for larger teachers while distillation reasoning gains persist.
  • The mid-training tradeoff is robust across teacher sizes, KL directions, and interpolation coefficients.

4 Why Does the Reasoning–Recall Tradeoff Occur?

The reasoning–recall tradeoff arises because teacher confidence varies across domains while students learn low-entropy facts earlier. High-entropy factual tokens therefore receive weaker KD supervision, motivating entropy-based routing.

  • Teacher supervision is asymmetric across data domains: Teacher predictive entropy is lower on procedural domains than on knowledge-intensive domains.This pattern is reported across OLMo-2 Instruct teacher sizes.
  • Teacher supervision is asymmetric across data domains: Within every domain and teacher size, lower entropy corresponds to higher teacher top-1 agreement with the ground-truth token.
  • Student knowledge state: By the end of pre-training, students acquire 67% of lowest-entropy facts but only 5% of highest-entropy facts.At mid-training initialization, unresolved facts are concentrated in the highest-entropy quintiles.
  • Knowledge distillation attenuates factual supervision: Higher teacher entropy lowers ground-truth probability and attenuates FKD and RKD ground-truth gradients to approximately 0.5× NTP for the highest-entropy facts.
  • Switch Distillation: Switch Distillation routes the lowest-entropy q% of tokens to reverse-KL distillation and trains the remainder with cross-entropy.Teacher entropy is computed per token at each training step.
  • Knowledge distillation attenuates factual supervision: During mid-training, both FKD and RKD produce their largest factual-recall deficits on high-entropy facts.

5 Switch Distillation Improves the Reasoning–Recall Tradeoff

Switch Distillation routes tokens using teacher predictive entropy, applying distillation to low-entropy predictions and next-token prediction otherwise. This targets the unequal reliability of teacher supervision across tokens while adding minimal computation.

  • Routing mechanism: Switch Distillation routes each token between distillation and next-token prediction using teacher predictive entropy.The lowest-q% entropy tokens are assigned to distillation, while the complement uses cross-entropy.
  • Routing mechanism: Low-entropy teacher predictions receive reverse-KL supervision because its mode-seeking behavior reinforces the teacher’s preferred continuation.
  • Efficiency: Teacher entropy is computed from logits already required for knowledge distillation, so Switch Distillation adds negligible computation and no parameters or forward passes.

6 Switch Distillation Experiments

Experiments show that Switch Distillation improves reasoning and knowledge performance while remaining closest to NTP on factual recall, and its reasoning advantages persist after post-training. Ablations attribute the gains primarily to entropy-based routing and reverse-KL supervision.

  • Mid-training results: 44.7%/42.1% reasoning macro-averages under Switch Distillation exceed NTP’s 26.1% with 7B and 13B teachers, respectively.Switch Distillation also records the strongest Knowledge & Commonsense performance, 49.3%/46.5%.
  • Mid-training results: 29.3%/29.3% Factual Recall under Switch Distillation remains closest among KD baselines to NTP’s 30.3%.
  • Post-training: 50.6% and 48.0% reasoning macro-averages after post-training preserve Switch Distillation’s lead from 7B and 13B teachers, respectively.Post-training modestly degrades Knowledge & Commonsense and Factual Recall, but Switch Distillation finishes with the highest Factual Recall macro-average.
  • Ablations: Replacing reverse KL with forward KL reduces Reasoning by 2.9% and Knowledge & Commonsense by 1.4%, with little Factual Recall change at -0.2%.
  • Ablations: Alternative routing signals underperform teacher entropy, while Teacher Top-1 Labels increase Factual Recall by 1.3% but reduce Reasoning by 6.4% and Knowledge & Commonsense by 2.8%.These results identify entropy-based routing and soft teacher distributions as important design choices.

7 Discussion

The paper argues that knowledge distillation should be stage-aware because teacher uncertainty and the student’s evolving knowledge interact differently during mid-training. Switch Distillation improves reasoning while largely preserving factual recall in a token-matched setting.

  • Discussion: Knowledge distillation should not be stage-agnostic because teacher uncertainty and student knowledge interact to produce a mid-training reasoning–recall tradeoff.
  • Discussion: Switch Distillation improves reasoning while largely preserving factual recall when training uses a token-matched setting.
  • Scope: The stage-aware principle may extend to late-stage or continual pre-training, but the paper studies mid-training.

A.1 Limitations

The study’s controlled experiments focus mainly on OLMo-2 because few model families release the checkpoints, data mixtures, and recipes needed for stage-dependent distillation analysis. The authors also identify selective-routing strategies and broader applicability as open directions.

  • Experimental scope: Controlled experiments focus on OLMo-2 because few model families provide complete checkpoints, data mixtures, and training recipes across stages.SmolLM2 and cross-family teacher analyses provide complementary evidence for generality.
  • Experimental choices: The experiments use relatively small 1B students, train beyond Chinchilla-optimal budgets, and do not exhaustively ablate these choices.The study prioritizes isolating stage- and objective-dependent distillation behavior rather than exploring every experimental dimension.
  • Method limitations: Richer learned routers may better identify when and where teacher supervision is beneficial than the entropy-based Switch Distillation router.The proposed router is cheap and effective but uses a simple primitive without additional supervision or parameters.
  • Scope of claims: Switch Distillation is developed and evaluated for mid-training, while benefits in late-stage pre-training remain a hypothesis for future work.The authors suggest testing it when factual acquisition slows under teacher supervision.

B.2 Ablating q

Routing-threshold ablations compare q values of 10%, 20%, and 30% across teacher settings. The authors select q = 20% because it maximizes reasoning while remaining competitive on factual recall and knowledge and commonsense.

  • Threshold selection: q = 20% provides the highest reasoning performance while remaining competitive or best on factual recall and knowledge & commonsense.This threshold is selected as the default across teacher sizes.

C.3 Analysis methodology

The analysis combines controlled token-level teacher-supervision measurements, gold-answer gradient comparisons, and robustness tests across training stages, teacher sizes, and model families. These analyses characterize why distillation’s reasoning–recall tradeoff changes during mid-training and motivate Switch Distillation.

  • Gold-answer gradients: The gold-answer gradient g measures the learning signal that a loss induces on accepted answer tokens, with larger values indicating stronger pressure toward the correct answer.The analysis derives g for NTP, forward-KL distillation, and reverse-KL distillation.
  • Gold-answer gradients: NTP pulls accepted-answer probability toward an absolute target of 1, whereas forward-KL supervision stops increasing it once the student matches the teacher’s probability mass.This contrast explains how teacher confidence bounds the factual learning signal.
  • Robustness across model families: In the SmolLM2 ecosystem, mid-training forward and reverse KL improve reasoning at the expense of factual recall, while Switch Distillation jointly improves both over NTP.The same qualitative stage dependence appears at substantially smaller scale, despite noisier curves.
  • Teacher supervision: Procedural tokens have lower teacher entropy than knowledge-intensive tokens, motivating entropy-based routing between distillation and next-token prediction.The entropy separation is evaluated across training stages and model families.
  • Teacher supervision: Across 12 stage-size cells, ROC AUC remains 0.744–0.826 with 95% CI width ≤0.005, so later training stages do not remove the entropy gap.The observed Base-to-Instruct erosion remains well above chance.
  • Objective comparison: FKL has higher CE–KL gradient alignment than RKL throughout training, and the widening gap at larger α indicates increasingly different update directions.This provides one explanation for the distinct positions of FKL and RKL on the reasoning–factual recall frontier.

D.5 Switch Distillation accelerates reasoning acquisition during mid-training

Switch Distillation accelerates reasoning acquisition during mid-training, surpassing the NTP baseline’s final reasoning level far earlier and continuing to improve throughout training.

  • D.5 Switch Distillation accelerates reasoning acquisition during mid-training: The trajectory comparison evaluates standard NTP, forward KD, and Switch Distillation across intermediate checkpoints in a 60B-token mid-training run.The analysis uses an OLMo-2 7B Instruct teacher and macro-averaged reasoning performance.
  • D.5 Switch Distillation accelerates reasoning acquisition during mid-training: Switch Distillation exceeds the NTP baseline’s final reasoning macro-average of 26.2% by the first evaluated checkpoint at 2.5B tokens.This corresponds to 1/24 of the NTP training token budget.
  • D.5 Switch Distillation accelerates reasoning acquisition during mid-training: Forward KD reaches the same 26.2% reasoning threshold after 5B tokens, twice the budget required by Switch Distillation.The comparison uses standard forward KD at α = 0.5 and the same 60B-token mid-training run.
  • D.5 Switch Distillation accelerates reasoning acquisition during mid-training: Both distillation methods continue improving with additional training and finish substantially above the NTP baseline.Switch Distillation maintains the strongest reasoning performance across the evaluated trajectory.

E.2 Ablation mid-training results

The ablation materials report per-task mid-training and intermediate post-training results, alongside analyses of teacher entropy, supervision alignment, and reasoning trajectories.

  • E.2 Ablation mid-training results: Figure 10 examines teacher entropy quintiles and reports that teacher entropy predicts factual acquisition under standard NTP.The same qualitative trend holds for OLMo-2 1B and 13B Instruct teachers.
  • E.2 Ablation mid-training results: Figure 11 analyzes KD on factual recall examples and reports the same qualitative trends for 1B and 13B Instruct teachers as for the 7B teacher.This analysis is presented as consistent with Figure 5.
  • E.2 Ablation mid-training results: Figure 12 reports that forward KL has higher gradient alignment with the cross-entropy objective than reverse KL throughout training.The figure compares the two KL directions through their gradient alignment with CE.
  • E.2 Ablation mid-training results: Figure 13 shows Switch Distillation surpassing the final reasoning performance of the 60B-token NTP baseline within the first 2.5B evaluated mid-training tokens.Standard forward KD reaches the same level after 5B tokens, and both methods continue improving over NTP.
  • E.2 Ablation mid-training results: Table 8 reports per-task results after intermediate post-training stages and uses NTP as the common reference for both teacher settings.Bold marks the best result within each teacher block, and asterisks indicate statistically significant improvements over the strongest competing baseline.
  • E.2 Ablation mid-training results: Table 9 reports per-task mid-training ablations of Switch Distillation using an OLMo-2 7B Instruct teacher.Macro-averaged results are reported separately in Table 3.
Loading 2609.01532v1…