Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall
- URL: http://arxiv.org/abs/2609.01532v1
- Date: Tue, 01 Sep 2026 17:00:30 GMT
- Title: Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall
- Abstract summary: We propose Switch Distillation, a simple mid-training objective that distills on tokens where the teacher is confident.<n>Relative to standard NTP, it achieves 1.61-1.71x the reasoning performance and 1.13-1.19x the knowledge and commonsense performance.
- Score: 79.59596652272923
- License: http://creativecommons.org/licenses/by/4.0/
- 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.
Related papers
- From Scoring to Acting: Outcome-Verified Comparative Self-Distillation for LLM Agents [16.67357658768852]
We propose Outcome-Verified Comparative Self-Distillation (OVCSD)<n>OVCSD organizes failed student rollouts into a prefix tree, adaptively invokes a skill-conditioned teacher from student-reached states, and retains only outcome-verified successful continuations.<n> Experiments on ALFWorld and WebShop across three model scales show that OVCSD consistently outperforms skill-free RL and existing self-distillation baselines.
arXiv Detail & Related papers (2026-07-30T09:47:58Z) - Prune-OPD: Efficient and Reliable On-Policy Distillation for Long-Horizon Reasoning [66.52232008796294]
Prune-OPD dynamically aligns training budgets with supervision quality.<n>It reduces training time by 37.6%--68.0% while preserving, and often improving, performance on challenging benchmarks.
arXiv Detail & Related papers (2026-05-08T14:38:53Z) - Distillation Traps and Guards: A Calibration Knob for LLM Distillability [54.90137955363471]
We propose a post-hoc calibration method that enables control over a teacher's distillability via reinforcement fine-tuning.<n>Our objective combines task utility, KL anchor, and across-tokenizer calibration reward.<n> Experiments show that students distilled from distillable teachers outperform SFT and KD baselines.
arXiv Detail & Related papers (2026-04-21T01:22:35Z) - Reinforcement-aware Knowledge Distillation for LLM Reasoning [63.53679456364683]
Reinforcement learning (RL) post-training has recently driven gains in long chain-of-thought reasoning large language models (LLMs)<n>Most existing knowledge distillation methods are designed for supervised fine-tuning (SFT), relying on fixed teacher traces or teacher-student Kullback-Leibler (KL) divergence-based regularization.<n>We propose RL-aware distillation (RLAD), which performs selective imitation during RL -- guiding the student toward the teacher only when it improves the current policy update.
arXiv Detail & Related papers (2026-02-26T00:20:39Z) - Explain in Your Own Words: Improving Reasoning via Token-Selective Dual Knowledge Distillation [4.517380184758914]
Token-Selective Dual Knowledge Distillation (TSD-KD) is a framework for student-centric distillation.<n>TSD-KD focuses on distilling important tokens for reasoning and encourages the student to explain reasoning in its own words.
arXiv Detail & Related papers (2026-02-25T09:58:49Z) - SGD-Based Knowledge Distillation with Bayesian Teachers: Theory and Guidelines [82.00660447875266]
Knowledge Distillation (KD) is a central paradigm for transferring knowledge from a large teacher network to a typically smaller student model, often by leveraging soft probabilistic outputs.<n>We rigorously analyze the convergence behavior of students trained with Gradient Descent (SGD)<n>Our analysis shows that learning from BCPs yields variance reduction and removes neighborhood terms in the convergence bounds compared to one-hot supervision.<n>Motivated by these insights, we advocate the use of Bayesian deep learning models, which typically provide improved estimates of the BCPs, as teachers in KD.
arXiv Detail & Related papers (2026-01-04T11:09:49Z) - Faithful Knowledge Distillation [75.59907631395849]
We focus on two crucial questions with regard to a teacher-student pair: (i) do the teacher and student disagree at points close to correctly classified dataset examples, and (ii) is the distilled student as confident as the teacher around dataset examples?
These are critical questions when considering the deployment of a smaller student network trained from a robust teacher within a safety-critical setting.
arXiv Detail & Related papers (2023-06-07T13:41:55Z) - Toward Student-Oriented Teacher Network Training For Knowledge Distillation [40.55715466657349]
We propose a teacher training method SoTeacher which incorporates Lipschitz regularization and consistency regularization into ERM.
Experiments on benchmark datasets using various knowledge distillation algorithms and teacher-student pairs confirm that SoTeacher can improve student accuracy consistently.
arXiv Detail & Related papers (2022-06-14T07:51:25Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.