Source-linked AI summary
Speculative Knowledge Distillation: Bridging the Teacher-Student Gap Through Interleaved Sampling
Wenda Xu, Rujun Han, Zifeng Wang, Long T. Le, Dhruv Madeka, Lei Li, William Yang Wang, Rishabh Agarwal, Chen-Yu Lee, Tomas Pfister
TL;DR
Large language model distillation must balance static-data distribution mismatch against low-quality student-generated samples that can mislead the teacher. SKD interleaves student proposals with teacher replacement of unlikely tokens, and it consistently outperforms existing KD methods across tasks, data sizes, and initialization strategies. The method also produces stronger speculative-decoding draft models, while on-policy KD can degrade when initialization is poor.
Problem
Supervised KD mismatches static training data with student-generated inference outputs, while on-policy KD can produce low-quality samples that yield inaccurate teacher feedback.
Method
SKD interleaves student token proposals with teacher replacement of tokens unlikely under the teacher, adapting between supervised-like and on-policy distillation.
Results
SKD consistently outperforms supervised and on-policy baselines across task types, data sizes, model families, and initialization strategies.
Takeaways & Limitations
SKD-trained students also improve speculative-decoding draft-model token acceptance from 71% to 85% higher and provide a 1.2X speed-up.
Takeaways & Limitations
On-policy KD can degrade over time when the initial student sample quality is insufficient, trapping the student in a sub-optimal state.
Abstract
from arXiv · showhide
Recent advances in knowledge distillation (KD) have enabled smaller student models to approach the performance of larger teacher models. However, popular methods such as supervised KD and on-policy KD, are adversely impacted by the knowledge gaps between teacher-student in practical scenarios. Supervised KD suffers from a distribution mismatch between training with a static dataset and inference over final student-generated outputs. Conversely, on-policy KD, which uses student-generated samples for training, can suffer from low-quality training examples with which teacher models are not familiar, resulting in inaccurate teacher feedback. To address these limitations, we introduce Speculative Knowledge Distillation (SKD), a novel approach that leverages cooperation between student and teacher models to generate high-quality training data on-the-fly while aligning with the student's inference-time distribution. In SKD, the student proposes tokens, and the teacher replaces poorly ranked ones based on its own distribution, transferring high-quality knowledge adaptively. We evaluate SKD on various text generation tasks, including translation, summarization, math, and instruction following, and show that SKD consistently outperforms existing KD methods across different domains, data sizes, and model initialization strategies.
1 INTRODUCTION
LLM compression is needed for practical deployment, but existing distillation methods face mismatches between training and inference distributions or low-quality student samples. SKD addresses these issues through adaptive teacher-student token interleaving and reports strong gains across tasks and settings.
- LLM compression is crucial for real-time deployment because large models impose substantial inference-time costs and memory demands.
- Supervised KD uses static outputs that mismatch student-generated inference samples, whereas on-policy KD can expose the teacher to low-quality, out-of-distribution student samples.On-policy KD is also highly sensitive to student initialization.
- SKD filters student-proposed tokens that the teacher considers unlikely and replaces them with teacher-generated tokens, gradually shifting from supervised-like to on-policy behavior.This interleaved design uses student samples while correcting low-quality intermediate tokens.
- SKD achieves substantial gains over supervised fine-tuning when distilling Gemma-7B to Gemma-2B, including 41.8% in translation, 230% in summarization, and 160% in arithmetic reasoning.Instruction-following training also yields 198% on MATH and 360% on GSMplus.
- The method combines adaptive token acceptance with broad robustness across task types, model initialization, data sizes, and model families.
2 PRELIMINARIES
The preliminaries define autoregressive token prediction and explain how temperature and top-k sampling shape generation. They also position SKD’s acceptance rule as inspired by, but distinct from, speculative decoding.
- An autoregressive language model predicts each next-token distribution conditioned on the input and previously generated prefix.Temperature is applied to logits before the softmax calculation.
- Higher sampling temperature produces more diverse outputs, while lower temperature yields more focused and deterministic generation.
- Top-k sampling restricts generation to the k highest-probability vocabulary tokens and motivates SKD’s teacher-based token acceptance criterion.
- Unlike speculative decoding, which ensures final samples follow the larger model’s distribution, SKD checks whether student tokens fall within the teacher’s top K tokens.
3 SPECULATIVE KNOWLEDGE DISTILLATION FOR LANGUAGE MODELS
SKD trains a student to imitate a teacher while reducing divergence on generated sequences. Its interleaved sampler accepts teacher-compatible student tokens and resamples rejected tokens from the teacher, spanning supervised and on-policy KD as limiting cases.
- Problem Setup: The distillation objective minimizes token-level divergence between fixed teacher and trainable student distributions on task sequences.Sequences may be generated by either model or supplied as ground truth.
- Baseline Knowledge Distillation Approaches: Supervised KD trains on fixed ground-truth data, while on-policy KD trains on student-generated samples using teacher and student token probabilities.ImitKD randomly mixes ground-truth and student-generated samples.
- Speculative Knowledge Distillation: SKD samples tokens from the student, checks each against the teacher’s top K tokens, and resamples rejected tokens from the teacher.The practical implementation proposes token blocks with γ = 5, although the algorithm illustration sets γ = 1.
- Speculative Knowledge Distillation: SKD’s acceptance rule preserves student samples likely under the teacher rather than replacing every token, avoiding degeneration into supervised KD.
- Speculative Knowledge Distillation: As student quality improves, SKD transitions from supervised-like correction toward on-policy interaction and degenerates to either baseline when all tokens are rejected or accepted.
4 EXPERIMENTAL SETUP
The experiments compare SKD with established distillation baselines across two model families, multiple generation tasks, and task-specific or task-agnostic data regimes. The setup varies data size and initialization while fixing key training hyperparameters and acceptance K.
- Student and Teacher Models: Experiments use Gemma and Qwen teachers and smaller students, including Gemma-7B-IT to Gemma-2B-IT and Qwen-7B-IT to Qwen-0.5B-IT.SFT student checkpoints are also used to test different initialization strategies.
- Baselines: SKD is compared with supervised fine-tuning, supervised KD, on-policy KD, and ImitKD; SeqKD is omitted because SFT consistently performed better in the authors’ comparison.
- Datasets and Tasks: Evaluation covers low-resource translation, dialogue summarization, arithmetic reasoning, and task-agnostic math instruction following.Task-specific experiments use 1K or 100 samples, while instruction-following experiments use 1K or 10K samples.
- Datasets and Tasks: Translation uses Flores-200 Assamese-to-English data and COMET evaluation, while dialogue summarization uses DialogSum and ROUGE-L.
- Datasets and Tasks: Arithmetic reasoning uses GSM8K answer accuracy, and math instruction following is evaluated on held-out GSMplus, MATH, Asdiv, and SVAMP sets.
5 RESULTS
SKD consistently outperforms baseline distillation methods across task-specific and task-agnostic generation, while remaining robust to initialization and low-data settings. Its gains arise from interleaved teacher-student sampling, which addresses low-quality student samples and training–inference mismatch.
- 5.1 TASK-SPECIFIC DISTILLATION: SKD outperforms all baseline KD approaches across three in-domain text-generation tasks and five of six test settings across two model families.The evaluated tasks include low-resource translation, dialogue summarization, and arithmetic reasoning.
- 5.2 TASK-AGNOSTIC DISTILLATION: SKD consistently outperforms supervised and on-policy KD across four held-out task-agnostic testing sets, with 20-40% gaps over baselines under 1k prompts.With 10k prompts, SKD can outperform or match all baseline methods; it also leads across seven math concepts.
- 5.3 DIFFERENT MODEL INITIALIZATION: SKD outperforms all baseline KDs with both instruction-tuned and supervised-fine-tuned student initialization, whereas on-policy KD can degrade from poor initialization.On-policy training may become trapped at a low-performance level when initial student samples are insufficiently high quality.
- 5.4 LOW DATA REGIME: At 100 data points per task, SKD outperforms supervised and on-policy KD across three tasks under both instruction-tuned and supervised-fine-tuned initialization.SFT overfitting in this regime can harm post-training KD, while end-to-end SKD bypasses the SFT stage.
- 5.6 SUMMARY AND ADDITIONAL RESULTS: SKD reduces estimated computation by 50% versus directly sampling from the teacher and improves speculative-decoding token acceptance by 71% to 85%, yielding a 1.2X speed-up.These additional results connect SKD’s training procedure to computational efficiency and downstream draft-model quality.
6 RELATED WORK
Knowledge distillation methods differ mainly in whether training data is fixed or generated on-policy, while SKD uses interleaved sampling to align student data with teacher guidance.
- Soft-label KD trains students by matching teacher and student output distributions, whereas hard-label KD mimics teacher input-output pairs.
- Supervised KD: Supervised KD uses fixed training data, creating a discrepancy between the student’s training and inference distributions.
- Speculative Knowledge Distillation: SKD assesses student-proposed tokens against the teacher’s top K tokens rather than sampling directly from the teacher distribution.
- Speculative Knowledge Distillation: Interleaved sampling is inspired by speculative decoding and is intended to produce high-quality on-policy samples for distillation.
7 CONCLUSION
SKD addresses limitations of existing knowledge distillation by interleaving student and teacher sampling. The method performs consistently across tasks, settings, model families, and dataset sizes, while also supporting speculative decoding.
- SKD uses interleaved sampling to reduce train-inference sample discrepancy and remove low-quality student-generated data.
- SKD consistently outperforms existing methods in both task-specific and task-agnostic distillation across model families, initializations, and dataset sizes.
- SKD-trained students outperform sequence-level mixing through token-level mixing and can accelerate speculative decoding.
8 REPRODUCIBILITY STATEMENT
The study documents reproducible implementation choices, evaluates a broad range of K values, and reports limitations of adaptive K strategies. Constant K=25 is used as the primary setting.
- Distillation objective: The implementation uses KL divergence between teacher and student token distributions as a knowledge-distillation objective.
- K-value study: A broad K range from 5 to 50 consistently outperforms supervised KD and on-policy KD across three text generation tasks.
- K-value study: The primary experiments use fixed K=25, while the authors note that task-specific K optimization could improve performance further.
- Adaptive K: Decreasing K during training did not outperform constant K, possibly because forced token replacement introduces off-policy behavior resembling supervised KD.
- Adaptive K: Adaptive decoding-time K outperformed most KD baselines across translation and summarization but did not outperform constant SKD with K=25.
- Future directions: The paper identifies combining Top-P and Top-K criteria as a direction for future SKD research.
D TRAINING AND HYPERPARAMETER DETAILS FOR BASELINES AND SKD
The experiments specify shared optimization settings, sampling controls, task-specific evaluation metrics, and training regimes for baseline and SKD comparisons across multiple tasks and model families.
- Optimization settings: Both instruction-tuned model families use learning rate 1e-5, disabled sampling dropout, task-specific sequence lengths, batch size 8, and gradient accumulation of 1.
- Sampling settings: The final sampling configuration uses temperature=0.5 and top-p=0.5 after unrestricted top-p sampling degraded generation quality, especially on math instruction.
- Teacher sampling: Teacher top-p and temperature affect resampled tokens, and lower teacher temperatures improved distillation in the reported grid search.
- Evaluation settings: Tables 5–7 compare baseline and SKD models across task-specific and task-agnostic settings, including 1,000-prompt and 10,000-prompt regimes.
- Evaluation metrics: Translation is evaluated with COMET, summarization with ROUGE-L, and mathematical tasks by numerical equivalence of parsed answers.
- Qualitative analysis: Qualitative samples show SKD avoiding hallucinated or repetitive translations and using teacher corrections to improve intermediate arithmetic reasoning.
I OVER-FITTING AT SFT STAGE CAN CAUSE SUB-OPTIMAL PERFORMANCE AT POST-TRAINING KDS
SFT overfitting can hurt post-training KD, while end-to-end SKD avoids dependence on a potentially suboptimal SFT checkpoint and also supports efficient speculative decoding.
- Over-fitting at SFT stage: Overfitting at SFT step 192 underperforms the converged step-128 checkpoint, while step 64 can outperform the overfitted checkpoint for both on-policy KD and SKD.This supports the claim that SFT-stage overfitting can produce suboptimal post-training KD initialization.
- Over-fitting at SFT stage: SKD bypasses the SFT stage through end-to-end training from an arbitrary base model, avoiding the checkpoint-quality constraint affecting post-training KD.
- Speculative decoding: SKD-trained students achieve the highest speculative-decoding speedup and token acceptance ratio among the evaluated draft models.SKD accelerates GEMMA-2B-IT by 1.21x in translation and 1.17x in summarization, while improving token acceptance ratio by 1.71x and 1.85x, respectively.
- Computational analysis: With a 90% expected acceptance rate and γ=5 in practice, SKD generates approximately 80% of tokens with the student and 20% with the teacher.The analysis estimates 4.1 student-generated tokens per run and rounds this to 4.
- Computational analysis: Under Nstudent=2B, Nteacher=7B, D=256, and L=200, SKD sampling costs approximately 1.72x on-policy KD but saves about 50% versus direct teacher sampling.Supervised KD is estimated at roughly 3.5x the cost of on-policy KD when ground-truth outputs must be sampled from the teacher.
N VALIDATION CONVERGENCE ACROSS METHODS
Across instruction-tuned and SFT student initializations, SKD reaches the lowest final validation loss across all three tasks, despite different early-training behavior among methods.
- Validation convergence: SKD achieves the lowest validation loss by the final iteration across all three tasks for both instruction-tuned and SFTed Gemma-2B students.The comparisons cover distillation from Gemma-7B over 350 training steps.
- Validation convergence: On-policy KD has higher validation loss early in training, whereas supervised KD can start lower on simpler tasks before SKD improves rapidly later.The early disadvantage of on-policy KD is attributed to lower-quality on-policy samples, while supervised KD risks overfitting later.
- Validation convergence: Improved student initialization enables on-policy KD to achieve lower final validation loss than supervised KD, indicating that on-policy methods benefit from better sample quality.
- Validation convergence: The methods remain stable under the KL-loss objective, so their convergence differences primarily reflect the effectiveness of their training samples.
O QUANTIFY SAMPLE OOD ISSUES WITH DIFFERENT KD APPROACHES
SKD reduces sample distribution mismatch in both directions: its samples are more familiar to the teacher than student proposals and more aligned with the student than ground-truth outputs.
- Teacher-side OOD: Teacher perplexity is 10.8 for SKD samples, versus 46.8 for student-proposed samples and 1.57 for ground-truth samples.This indicates that the teacher assigns substantially lower perplexity to SKD samples than to student proposals.
- Evaluation setup: The OOD analysis compares ground-truth outputs, student on-policy samples, and SKD interleaved-sampling outputs using teacher and student per-sample perplexity.The experiment uses 100 randomly selected translation prompts for student and SKD sampling.
- Student-side OOD: Student perplexity is 358 for SKD samples, versus 5606.3 for ground-truth samples, showing stronger alignment with the student distribution.Ground-truth samples therefore remain substantially more off-policy for the student than SKD-generated samples.