Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall

Hugging Face Daily Papers Papers
knowledge-distillation mid-training reasoning factual-recall language-models switch-distillation

Summary

This paper introduces Switch Distillation, a novel mid-training objective that selectively applies knowledge distillation based on teacher confidence to improve reasoning and preserve factual recall in smaller language models.

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.
Original Article
View Cached Full Text

Cached at: 09/02/26, 11:48 PM

Paper page - Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall

Source: https://huggingface.co/papers/2609.01532 Authors:

,

,

,

,

,

,

,

,

,

,

Abstract

Switch Distillation selectively applies logit-based knowledge distillation during mid-training based on teacher confidence to improve reasoning and preserve factual recall in smaller language models.

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--withpost-trained teachersbehaves fundamentally differently duringmid-training, an intermediate phase ofself-supervised learningon curated corpora. Surprisingly, while forward KD simultaneously improvesreasoningandfactual recallduring pre-training relative to standardnext-token prediction(NTP), it instead slowsfactual recallacquisition duringmid-trainingdespite continuedreasoninggains. 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 proposeSwitch Distillation, a simplemid-trainingobjective that distills on tokens where the teacher is confident, using teacherpredictive entropyas a lightweight routing signal, and otherwise falls back tocross-entropy.Switch Distillationconsistently outperforms existing distillation objectives across teacher sizes. Relative to standard NTP, it achieves 1.61-1.71x thereasoningperformance and 1.13-1.19x the knowledge and commonsense performance while preserving 96.7-96.8% offactual recall. Crucially, these benefits persist after post-training:Switch Distillationcloses thefactual recallgap while maintaining 1.25-1.32x and 1.13-1.20x gains inreasoningand knowledge and commonsense, respectively.

View arXiv pageView PDFGitHubAdd to collection

Get this paper in your agent:

hf papers read 2609\.01532

Don’t have the latest CLI?curl \-LsSf https://hf\.co/cli/install\.sh \| bash

Models citing this paper0

No model linking this paper

Cite arxiv.org/abs/2609.01532 in a model README.md to link it from this page.

Datasets citing this paper0

No dataset linking this paper

Cite arxiv.org/abs/2609.01532 in a dataset README.md to link it from this page.

Spaces citing this paper0

No Space linking this paper

Cite arxiv.org/abs/2609.01532 in a Space README.md to link it from this page.

Collections including this paper0

No Collection including this paper

Add this paper to acollectionto link it from this page.

Similar Articles

Improving Reasoning Capabilities in Small Models through Mixture-of-Layers Distillation with Stepwise Attention on Key Information

arXiv cs.CL

This paper proposes a novel Chain-of-Thought distillation framework that transfers teacher models' stepwise attention on key information to student models through a Mixture-of-Layers module for dynamic layer alignment. The method achieves consistent performance improvements on mathematical and commonsense reasoning benchmarks by explicitly guiding student models to progressively focus on critical information during reasoning.

Masked Distillation: Internalizing the Chain-of-Thought in Language Models

arXiv cs.AI

Masked distillation is a knowledge-distillation framework that trains a student LLM to predict only solution tokens while a reasoning teacher provides feedback, aiming to internalize chain-of-thought computation into model parameters. The method shows task-dependent success, working on GSM8K but requiring small scaffolds for harder tasks like Countdown.

Reasoning Compression with Mixed-Policy Distillation

arXiv cs.AI

This paper proposes Mixed-Policy Distillation (MPD), a framework that transfers concise reasoning behaviors from large teacher models to smaller student models, reducing token usage by up to 27.1% while improving performance.