Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall
Knowledge distillation during mid-training enhances reasoning but slows factual recall; proposes Switch Distillation.
Key Findings
Methodology
The study uses forward KL distillation and Switch Distillation to analyze asymmetric teacher confidence during mid-training. Switch Distillation routes distillation based on teacher predictive entropy, ensuring distillation when teacher confidence is high, otherwise falling back to cross-entropy.
Key Results
- Switch Distillation achieves 1.61-1.71x reasoning performance and 1.13-1.19x knowledge and commonsense performance over standard NTP, while preserving 96.7-96.8% factual recall.
- Post-training, Switch Distillation maintains 1.25-1.32x reasoning and 1.13-1.20x knowledge gains.
- Experiments show Switch Distillation consistently outperforms existing distillation objectives across teacher sizes.
Significance
The study reveals unique behavior of knowledge distillation during mid-training and proposes Switch Distillation to effectively mitigate the reasoning-factual recall tradeoff. This finding provides new insights into training strategies for language models, especially in data-scarce scenarios, enhancing reasoning capabilities.
Technical Contribution
Technical contributions include proposing Switch Distillation, which uses teacher predictive entropy for routing, significantly improving reasoning performance during mid-training while maintaining factual recall integrity. This method offers new perspectives for knowledge distillation during mid-training.
Novelty
This is the first to reveal the differential impact of knowledge distillation on reasoning and factual recall during mid-training, proposing Switch Distillation as a solution with significant advantages over traditional methods.
Limitations
- Switch Distillation may not fully utilize teacher guidance in high-entropy scenarios, leading to slight factual recall decline.
- The method's performance on specific domain data may be limited by teacher predictive entropy.
- Effectiveness may be less than expected in extreme data scarcity scenarios.
Future Work
Future research could explore applying Switch Distillation to a broader range of datasets and tasks to further optimize its performance across different training stages. Additionally, investigating how to combine other distillation strategies to enhance overall model capabilities is a promising direction.
AI Executive Summary
In training language models, knowledge distillation is a common method where smaller student models learn from stronger teacher models. However, its effectiveness during mid-training has been debated. This paper finds that traditional forward KL distillation enhances reasoning but slows factual recall acquisition during mid-training.
To address this, the researchers propose Switch Distillation, which uses teacher predictive entropy to decide whether to distill. Specifically, distillation occurs when the teacher is confident about a data point; otherwise, it falls back to cross-entropy supervision. This method excels in experiments, outperforming traditional methods in reasoning and knowledge acquisition.
Switch Distillation not only offers a new solution for knowledge distillation during mid-training but also points to future research directions. Although the method may face challenges in certain scenarios, its ability to enhance reasoning while maintaining factual recall shows great potential.
Deep Analysis
Background
In recent years, training language models has increasingly relied on knowledge distillation, where smaller student models learn from stronger teacher models. However, the effects of knowledge distillation across different training stages have not been fully explored, especially during mid-training. Mid-training typically involves self-supervised learning on high-quality datasets to enhance model reasoning and factuality.
Core Problem
The core problem is the inconsistent impact of traditional knowledge distillation methods on reasoning and factual recall during mid-training. Specifically, while reasoning improves, factual recall acquisition slows. This inconsistency may stem from asymmetric teacher confidence across data domains.
Innovation
The core innovation is the proposal of Switch Distillation, which routes distillation based on teacher predictive entropy. Unlike traditional methods, this approach distills when teacher confidence is high, otherwise falling back to cross-entropy supervision, effectively mitigating the reasoning-factual recall tradeoff.
Methodology
- �� Conduct initial experiments using forward KL distillation to analyze its performance during mid-training.
- �� Propose Switch Distillation, routing distillation based on teacher predictive entropy.
- �� Distill when teacher confidence is high, otherwise fall back to cross-entropy supervision.
- �� Conduct experiments across multiple teacher sizes to validate the method's effectiveness.
Experiments
Experiments use the OLMo-2 ecosystem, including 1B student models and 7B, 13B teacher models. Datasets include Dolmino Mix 1124, with experiments covering different distillation strengths and KL directions. Performance of Switch Distillation is evaluated against standard NTP and other distillation methods.
Results
Results show Switch Distillation achieves 1.61-1.71x reasoning performance and 1.13-1.19x knowledge and commonsense performance over standard NTP, while preserving 96.7-96.8% factual recall. These improvements persist post-training, demonstrating the method's robustness.
Applications
Switch Distillation can be applied in scenarios requiring high reasoning capabilities, such as intelligent assistants and automated question-answering systems. Its advantages in data-scarce situations make it suitable for resource-limited applications.
Limitations & Outlook
While Switch Distillation excels in reasoning, it may not fully utilize teacher guidance in high-entropy scenarios, leading to slight factual recall decline. Additionally, its performance on specific domain data may be limited by teacher predictive entropy.
Plain Language Accessible to non-experts
Imagine you're cooking in a kitchen. Traditional knowledge distillation is like following a recipe, where the recipe tells you what to do at each step. But during mid-training, this recipe sometimes isn't quite right because it lacks confidence in certain steps. Switch Distillation is like a smart assistant that decides whether to follow the recipe based on the difficulty and importance of each step. If a step is simple, it lets you do it yourself; if it's complex, it gives you more guidance. This way, you not only make delicious dishes but also learn more cooking skills.
ELI14 Explained like you're 14
Hey there! Imagine you're playing a super complex game, and there's a super smart NPC (non-player character) teaching you how to level up. Usually, this NPC tells you what to do at each step, but sometimes it makes mistakes, especially in tough levels. Switch Distillation is like a smart game assistant that decides whether to let the NPC guide you based on the difficulty of each level. If a level is easy, it lets you explore on your own; if it's hard, it gives you more hints. This way, you not only beat the game but also become a game master!
Glossary
Knowledge Distillation
A method where smaller student models learn from stronger teacher models to enhance performance.
Used in this paper to enhance reasoning capabilities of language models.
Forward KL Distillation
A method that distills by minimizing the forward KL divergence between teacher and student predictive distributions.
Used in experiments to analyze mid-training effects.
Switch Distillation
A new method using teacher predictive entropy to route the distillation process, aiming to enhance reasoning performance during mid-training.
The core method proposed in this paper to address the reasoning-factual recall tradeoff.
Predictive Entropy
Measures the confidence of a model's prediction; lower entropy indicates higher confidence.
Used to decide whether to distill.
Dolmino Mix 1124
A dataset comprising multiple data sources, used for training and evaluating language models.
Used in experiments to test Switch Distillation performance.
Open Questions Unanswered questions from this research
- 1 How to apply Switch Distillation to more diverse datasets to verify its generality.
- 2 Exploring the possibility of combining other distillation strategies to further enhance model performance.
- 3 Investigating the applicability of Switch Distillation in extreme data scarcity scenarios.
Applications
Immediate Applications
Intelligent Assistants
Enhance reasoning capabilities of intelligent assistants using Switch Distillation for better performance in complex dialogues.
Automated Q&A Systems
Improve accuracy and efficiency of Q&A systems in handling complex questions, especially in data-limited scenarios.
Long-term Vision
General AI
Advance general AI development by optimizing knowledge distillation during mid-training, despite challenges in data and computational resources.
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.