Speculative Knowledge Distillation: Bridging the Teacher-Student Gap Through Interleaved Sampling

## Abstract

Abstract 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

Text generation capabilities of large language models (LLMs) have seen continuous improvement, largely driven by scaling up the number of parameters and the amount of training data (Kaplan et al., ; Anil et al., ; Dubey et al., ; OpenAI et al., ). However, the substantial inference-time costs and memory footprint associated with LLMs present significant challenges for practical deployment (Agarwal et al., ). Therefore, compressing LLMs while maintaining their performance is crucial for real-time practical applications.

Knowledge Distillation (KD) (Hinton et al., ) is a widely used method to compress LLMs by transferring knowledge from a larger teacher model to a smaller student model. Traditional KD approaches, such as supervised KD (Sanh et al., ) and SeqKD (Kim & Rush, ), rely on a static dataset of outputs to train the student model. However, this fixed dataset can lead to a distribution mismatch between the training data and the student’s generated samples at inference time, hindering the student’s learning. On-policy KD (Agarwal et al., ) attempts to address this train-inference mismatch by training the student on its self-generated samples, with the teacher provides feedback through its token-level probabilities. However, the student may generate low-quality samples that are out-of-distribution (OOD) for the teacher, especially early in the training, leading to inaccurate teacher feedback. Additionally, the performance of on-policy is highly sensitive to the student’s initialization, as illustrated in Figure 4.

In this work, we introduce Speculative Knowledge Distillation (SKD) that uses interleaved sampling between teacher and student to overcome the challenges faced by supervised and on-policy distillation (Figure 2). Similar to on-policy KD, SKD utilizes student-generated samples to address the train-inference mismatch. However, to mitigate the issue of low-quality student samples, SKD filters out intermediate tokens that the teacher is unlikely to produce and instead re-samples them from the teacher, drawing inspiration from speculative decoding (Leviathan et al., ). The design of SKD is supported by theoretical results in imitation learning (Ross & Bagnell, ; Ross et al., ), which suggest that gradually shifting the sample generation from the teacher to the student is an effective approach. Early in training, when the student model is likely to propose many out-of-distribution tokens, SKD functions similarly to supervised KD, replacing low-quality tokens with those from the teacher. As training progresses and the student’s sample quality improves, SKD behaves more like on-policy KD, accepting many of the student’s proposed tokens.

Figure: Figure 2: Overview of Speculative Knowledge Distillation (SKD) on an arithmetic reasoning task. Left: SKD addresses the limitations of on-policy knowledge distillation (KD) by filtering out low-quality student samples and replacing them with teacher generated tokens. Right: illustration of how SKD generalizes to both supervised KD (replacing with all teacher tokens) and on-policy (accepting all student tokens).
Refer to caption: x2.png

SKD significantly outperforms supervised and on-policy distillation for LLMs, as shown in Figure 1. Specifically, when distilling Gemma-7B to Gemma-2B, SKD achieves substantial gains compared to supervised fine-tuning in low-resource machine translation (41.8%), summarization (230%), and arithmetic reasoning (160%). Furthermore, training SKD on instruction-following data leads to impressive results on held-out MATH (Hendrycks et al., ) (198%) and GSMplus (Li et al., ) (360%) testing suites. To provide a comprehensive understanding of SKD’s capabilities, we conduct a systematic evaluation under varying conditions and offer practical guidelines for tokens acceptance criteria from the teacher model, considering factors such as task type and model initialization. Our key contributions are:

- •
A novel KD approach that utilizes on-policy student samples likely to be generated by the teacher, mitigating the issues of low-quality samples in on-policy KD and dynamically switching between supervised and on-policy KD.
- •
Superior performance across various text generation tasks, including summarization, translation, and reasoning tasks, in both task-specific and task-agnostic scenarios.
- •
Robustness and adaptability across different model initialization, data sizes, and model families.

## 2 Preliminaries

#### Language Models.

An auto-regressive language model, denoted as $M$, predicts the probability distribution of the next token in a sequence. Conditioned on an input sequence $X$ and the generated prefix $y_{<i}=\{y_{1},y_{2},...,y_{i-1}\}$ consisting of $(i-1)$ tokens from a vocabulary $V$, $M$ calculates the probability of the $i$-th token, $p_{y_{i}}=M(y_{i}|X,y_{<i})$.
The auto-regressive language model M computes a logit $l$ at the $i$-th token position. To obtain the probability $p_{y_{i}}$, we apply the softmax function with temperature $t$: $p_{y_{i}}=\frac{exp(l_{i}/t)}{\sum_{j=1}^{|V|}exp(l_{i}/t)}$. The temperature $t$ introduces randomness into the text generation process. A higher temperature leads to more diverse sample output, while a lower temperature results in more focused and deterministic generation. In practical implementations, the vocabulary $V$ is often truncated using top-k sampling (Radford et al., ),
which improves text generation compared to searching over the entire vocabulary distribution. We draw inspiration from top-k sampling to develop our acceptance criteria for student-generated tokens. We include background of divergence metrics in the Appendix A.

#### Speculative Decoding (Leviathan et al., 2022 ; Chen et al., 2023 )

accelerates the decoding process by employing a smaller draft model to generate token candidates, which are then verified by larger approximation models. Our interleaved sampling is inspired by speculative decoding but introduces key differences. While speculative decoding ensures that final sampled tokens adhere to the larger model’s distribution, our method directly assesses the feasibility of student-proposed tokens within the teacher’s top $K$ tokens (See Sec 3.2).

## 3 Speculative Knowledge Distillation for Language Models

#### Problem Setup.

Given a teacher LM $M_{t}$ and a student LM $M_{s}$, with differing model capacities, our goal is to train $M_{s}$ to mimic the behavior of $M_{t}$ on a specific task $T$. We assume that $M_{s}$ has learnable parameters $\theta_{s}$, while $M_{t}$’s parameters $\theta_{t}$ are fixed. For task $T$, we have a set of prompts $\{x\}$ and corresponding output sequences $\{y\}$, which can be generated by either model or provided as ground truth. To measure the divergence between the token-level distributions of $M_{t}$ and $M_{s}$ for a given input-output pair $(x,y)$, we define the following metric:

$$ $D(M_{t}||M_{s})(y|x)=\frac{1}{L_{y}}\sum_{i=1}^{L_{y}}D(M_{T}(.|y_{<i},x)||M_{ s}(.|y_{<i},x))$ (1) $$

where $L_{y}$ is the length of the output sequence $y$, and $i$ indicates decoding steps. Our training objective is to minimize $D(M_{t}||M_{s})(y|x)$ such that $M_{s}$ can effectively imitate $M_{t}$ on task $T$.

### 3.1 Baseline Knowledge Distillation Approaches

#### Supervised FT

is a training method where a student policy is trained on a fixed dataset of input-output pairs ($x,y$). The objective is to minimize the negative log-likelihood of the student’s predictions: $L_{SFT}=-\log p_{s}(y|x)$. It’s also a common initialization strategy for more advanced knowledge distillation techniques like Supervised KD, On-Policy KD, and SKD.

#### SeqKD (Kim & Rush, 2016a )

is a supervised fine-tuning method that trains on sequences generated by a teacher model using maximum likelihood.

#### Supervised KD (Sanh et al., 2020 )

trains a student model to mimic the token-level probability distribution of a teacher model over a fixed ground truth dataset. This is achieved by minimizing the distance between the student’s and teacher’s predictions, as defined in Equation 1. The ground truth labels, denoted by y, are obtained from the fixed dataset.

#### On-Policy KD (Agarwal et al., 2024 )

addresses the training-inference mismatch in student models by sampling target tokens directly from the student’s own output. This approach calculates both the teacher and student’s token probabilities and employs Equation 1 to train the student model on its self-generated data.

#### ImitKD (Lin et al., 2020 )

is a hybrid approach that can be viewed as naively combining supervised and on-policy KD. During training, it randomly selects samples from either the ground truth dataset or the student model’s own output with equal probability. This mixed dataset is then used to train the student model, optimizing the distance metric defined in Equation 1.

### 3.2 Speculative Knowledge Distillation

Speculative knowledge distillation (SKD) draws two important inspirations from imitation learning (Ross et al., ; Ross & Bagnell, ): 1) Directly taking samples from a fixed dataset or sampling from teacher can cause the state distribution to be unrealistically good. As a result, student model never learns to correct previous mistakes, leading to poor performance at test time. 2) Directly drawing samples from student can result in low quality samples that can be out-of-distribution for the teacher, resulting in inaccurate feedback when assigning logits to the student.
SKD addresses above limitations by using major components from speculative sampling (Leviathan et al., ): 1) student model proposes on-policy sample candidate, and 2) the teacher model evaluates these proposed samples concurrently and accepts only those that it deems likely to generate. After speculative sampling, SKD calculates teacher and student’s token probabilities and employ equation 1 to train the student model (L9 in Algorithm 1).

Figure: Algorithm 1 Speculative knowledge distillation

#### Interleaved On-Policy Sampling

Given a prompt $x$, student model $M_{s}$ will generate a sequence of tokens, $(y_{i},...,y_{i+\gamma})$, following its own distribution $M_{s}(.|y_{<i},x)$ (L4 in Algorithm 1). The teacher model $M_{t}$ will evaluate this sequence, accepting or rejecting each token based on pre-defined criteria, which resembles speculative decoding (Leviathan et al., ). If a token $y_{i}$ is rejected, the subsequent tokens ($y_{i+1}$, …) will be discarded, and $y_{i}$ will be resampled from $M_{t}$’s distribution $M_{t}(.|y_{<i},x)$ (L5-7 in Algorithm 1). The student model will then continue to generate tokens on-policy based on $M_{s}(.|y_{<i},x,y_{i})$ for the remaining $\gamma$ tokens. This sample generation process is particularly advantageous for auto-regressive models, where early errors can propagate and affect future predictions (Ross & Bagnell, ). By focusing on on-policy samples that are likely to be generated by the teacher model, we can effectively utilize limited student capacity to mimic teacher.

#### Token Acceptance Criteria

Using the default acceptance criteria in speculative decoding corresponds to sampling from teacher’s distribution, which would degenerate to supervised KD. Instead, we are interested in generating student samples that are likely under the teacher.
To do so, we draw inspiration from top-k sampling (Radford et al., ) that draws samples using only the k highest probability tokens from the entire vocabulary $V$ at each generation step.
Following this intuition, we define our acceptance criteria based whether the student token $y_{i}$ falls into the top $K$ tokens of teacher’s distribution $M_{t}(.|y_{<i},x)$. Similar to speculative decoding, the teacher $M_{t}$ can evaluate sequence $(y_{i},...,y_{i+\gamma})$ proposed by student in parallel and return to student with the token position that discarded.
We replace the discarded student token by re-sampling a token from the teacher’s distribution. Note that in Algorithm 1, we simplify the SKD presentation by setting $\gamma=1$, but in practice use $\gamma=5$ for efficient implementation.

#### Sample Transition During SKD Training.

In the early stages of training, student models often propose low-quality samples that are frequently rejected by the teacher. This initial behavior resembles supervised KD, where the student is explicitly corrected and guided towards the teacher’s generated samples. As the student’s training progresses, the quality of its proposed samples improves, leading to a more on-policy interaction with the teacher. SKD dynamically adjusts the balance between supervised and on-policy KD based on the distribution gap between the teacher and student models. This adaptive approach allows the student to efficiently learn from the teacher’s samples. Moreover, SKD is a flexible framework that can degenerate both supervised and on-policy KD as special cases. In Figure 2, when the student’s samples are consistently rejected or accepted by the teacher, SKD effectively degenerates to supervised or on-policy KD, respectively. This demonstrates the generality and versatility of SKD as a unified approach to knowledge distillation.

## 4 Experimental Setup

Student and Teacher Models. We conducted experiments using two model families, Gemma1 (Team et al., ) and Qwen2 (Yang et al., ). We use supervised fine-tuned (SFT) Gemma-7B-it (it indicates instruction-tuned) and Qwen-7B-it as the teacher models, while Gemma-2B-it and Qwen-0.5B-it were employed as student models. We further leveraged SFTed student models to compare KDs under different model initializations.

Hyperparameters. For all fine-tuning process, we use learning rate $1$e-$5$, warmup ratio $0.1$ and $0.1$ drop out rate for all fine-tuning processes. All SFT checkpoint is trained for three epoches and we select the checkpoint with the lowest validation loss. Details about SFT training dataset can be found in Section 4.1. Without parameter search, we fixed acceptance criteria $K$ to be $25$ throughout main experiment and ablation ($K$=$25$ is a common choice for top-k sampling (Radford et al., ). See Appendix B for detailed study). We include training details and our hyper-parameter choice of baselines and SKD in the Appendix Section C.

#### Baselines.

We conduct SKD over widely accepted KD approaches listed in Section 3.1: Supervised Fine-tuning, Supervised KD, on-policy KD and ImitKD. We compared SeqKD to supervised fine-tuning (SFT) using ground truth data and found that SFT consistently outperformed SeqKD. Therefore, we did not include SeqKD results.

### 4.1 Dataset and Evaluation.

We evaluated SKD and baseline methods across three generation tasks: low-resource translation, dialogue summarization, and arithmetic reasoning. Task-specific KD often operates with training data around 1K samples (Quteineh et al., ). Thus, we randomly draw 1K samples from our task data, and further test low-data regime with 100 samples. Additionally, we performed experiments using 1K and 10K on a task-agnostic instruction-following task in the math domain. It’s important to note that different baseline KD methods have access to different training data. Supervised KD, SFT, and ImitKD have access to both $(x,y)$ pairs from the ground truth dataset, while on-policy and SKD only see the prompt $x$ and self-generate the target label during training. Since our assumption that the teacher model should be an expert in a specific task, we used full training data for each task to supervise the FT teacher model. For SFTed student initialization, we supervise the FT student model with our constructed training set $x,y$ (around 1000 instances). For reproducibility, all baselines and SKD use greedy decoding for evaluation. We include all results of supervised FT teacher model, supervised FT student model and instruction tuned student models in the Appendix D.

#### Low Resource Translation.

We utilize the Flores-200 (Team, ) Assamese-to-English translation dataset for low-resource translation. Due to the scarcity of Assamese-to-English training data, we employ the Flores-200 development set (997 instances) as our training set. Additionally, we split the Flores-200 testing set (1012 instances) into a development set (500 instances) and a testing set (512 instances) for evaluation. We use SOTA learned metric COMET (Rei et al., ) to evaluate the translation quality. We use 997 instances from our constructed training set to SFTed both teacher and student models.

#### Dialogue Summarization.

For dialogue summarization, we utilize the DialogSum dataset (Chen et al., ). We randomly sample 1K instances from the DialogSum training set to create our input-output pairs $(x,y)$. For evaluation, we employ the dataset’s development set (500 instances) and test set (1500 instances). To assess summary quality, we use the ROUGE-L metric (Lin, ). We Supervised FT our teacher models on 10K instances from the DialogSum training set, and our student models on the 1K instances from our constructed training set.

#### Arithmetic Reasoning.

For arithmetic reasoning, we utilize the GSM8K dataset (Cobbe et al., ). We split the GSM8K training set into a development set (473 instances) and a training set (7K instances). We randomly sample 1K instances from the training set to create our input-output pairs $(x,y)$. For evaluation, we employ the development set and the GSM8K test set (1319 instances). We assess the quality of the generated output using answer accuracy. We Supervised FT our teacher models on 7K instances from the GSM8K training set, and our student models on the 1K instances from our constructed training set.

#### Math Instruction Following.

For math instruction following, we utilize the UltraInteract dataset (Yuan et al., ). We randomly sample 11K instances from the UltraInteract training set, dividing them into a training set (10K instances) and a development set (1K instances). For evaluation, we employ the GSMplus (Li et al., ), Math (Hendrycks et al., ), Asdiv (Miao et al., ) and SVAMP (Patel et al., ) test sets as held-out sets. We Supervised FT our teacher models on 10K instances from the UltraInteract training set, and our student models on the 1K instances from our constructed training set. We also tested the setting with 10K training instances for student model.

We explored two types of knowledge distillation: task-specific (Quteineh et al., ) and task-agnostic (Agarwal et al., ). In task-specific distillation, training and inference data share the same format, and the task is known during training. In task-agnostic distillation, only the domain is known, while the specific task varies during deployment. We categorized distillation of translation, dialogue summarization, and arithmetic reasoning as task-specific and distillation of math instruction following as task-agnostic. For both math instruction following and arithmetic reasoning, we use chain-of-thought (Wei et al., ) prompts to generate answer. All prompt formats can be found in Appendix E. Details about parsing and evaluation for all tasks can be found in the Appendix F.

## 5 Results

We validate SKD’s effectiveness through comprehensive studies, including task-specific (Section 5.1) and task-agnostic distillation (Section 5.2), different model initializations (Section 4), low-data regimes (Section 5.4), and an ablation study on two-staged training (supervised KD followed by on-policy KD) (Section 5.5).

**Table 1: SKD outperforms baseline KDs at in-domain tasks: low resource Assamese-to-English translation, dialogue summarization, arithmetic reasoning. We use Gemma-2b and Qwen-0.5B instruction tuned model (w/o SFT) as students and SFTed Gemma-7B-it and Qwen2-7B-it as teachers. $K$=$25$ for SKD.**
|  | Supervised FT Gemma-7B to Gemma-2B | Supervised FT Qwen2-7B to Qwen2-0.5B |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | Translation | Summarization | GSM8K | Translation | Summarization | GSM8K |
| Baselines | COMET | ROUGE-L | Acc | COMET | ROUGE-L | Acc |
| SFT | 72.5 | 31.7 | 18.7 | 57.6 | 29.2 | 31.8 |
| Supervised KD | 73.3 | 34.1 | 22.5 | 57.5 | 31.3 | 35.0 |
| On-policy KD | 36.1 | 34.1 | 25.3 | 53.0 | 31.2 | 33.9 |
| ImitKD | 74.8 | 34.9 | 26.2 | 55.9 | 30.9 | 33.2 |
| SKD | 75.3 | 35.0 | 29.1 | 57.1 | 31.7 | 36.6 |

### 5.1 Task-specific distillation

In this section, we investigated task-specific distillation, where the training and inference prompt-response sets share the same format, and the task is known during training. We study two model families, supervised FT Gemma-7B to instruction tuned Gemma-2B (Gemma-2B-it) and supervised FT Qwen2-7B to instruction tuned Qwen2-0.5B (Qwen2-0.5B-it). As shown in Table 1, the performance of supervised KD and on-policy KD is highly task-dependent, likely due to the model’s initial familiarity with each task (see Sec 4 for details). SKD, trained on Gemma-2B-it, consistently outperformed all baseline KD approaches across three in-domain text generation tasks.

When initialized with Qwen2-0.5B-it, SKD outperformed all baselines on summarization and GSM8k. We found that on-policy KD consistently underperformed supervised KD and SKD for translation and summarization. This is primarily attributed to the low-quality samples generated by the student model during training. Although ImitKD (partially on-policy) can achieve competitive performance with a stronger student like Gemma-2B due to its mixed data strategy, it falls behind SKD when using the Qwen-0.5B student due to the lower quality of on-policy samples (See example in Sec 5.3). Our results demonstrate SKD’s superiority in five out of six test settings across two model families, highlighting its optimized performance for task-specific distillation.

Figure: Figure 3: SKD outperforms baseline KDs at four held-out testing sets (Math, GSMplus, SVAMP and ASDiv). We consider two data size setting, 1k and 10k respectively and report testing accuracy in the first row. In the second row, we calculate the performance gain of SKD over SFT and report the relative performance gains of baselines compared to SKD. As SKD outperforms all baselines in most cases, these values typically range from 0 to 1. In some instances, supervised KD may be outperformed by SFT, resulting in negative values. In the middle figure at the second row, we show that SKD outperforms all baselines under $7$ math concepts at math testing set. SKD is performed under $K$=$25$.
Refer to caption: x3.png

### 5.2 Task-agnostic distillation

To explore task-agnostic instruction following, where the exact nature of the task is unknown during training and can vary during deployment, we trained and evaluated SKD and baselines on a math domain instruction following task. Our training set included diverse math instruction following data, such as MathQA (Amini et al., ), math reasoning (Mishra et al., ), and tabular processing (Lu et al., ). We tested each KD approach’s performance on four held-out testing sets, ranging from simple math word problems (ASDiv and SVAMP) to grade-level math problems (GSMplus) and competitive math problems (Math).

As illustrated in Figure 3, SKD initialized with Gemma-2B-it consistently outperformed supervised KD and on-policy KD across these four held-out sets when trained on 1k prompts. We also calculated the performance gain of SKD over SFT and reported baseline performance gains relative to SKD. Results showed that under 1k prompts, supervised KD and on-policy could approach close performance on simple tasks like SVAMP and ASDiv but still had a 20-40% performance gap compared to SKD (Second row on the left of Figure 3). The middle figure demonstrates SKD’s superiority over baselines on the Math task under 7 different concepts. Since the training set was task-agnostic instruction following data, it could be trained on a larger scale in practical scenarios. On the right side of Figure 3, we show that SKD can outperform or match all baseline methods with 10k prompts. Our results demonstrate SKD’s generalization capability under task-agnostic distillation. Detailed table result can also be found at Appendix D.

### 5.3 Different model initialization

In this section, we explore how different student model initialization impact the effectiveness of various KD approaches. As shown in Figure 4, on-policy KD initialized with an instruction-tuned student model can surprisingly degrade performance over time. This raises the question of whether a teacher’s token-wise probability distribution always provides beneficial guidance. Our findings suggest that if the initial student sample quality is insufficient, the student model may become trapped in a sub-optimal state, unable to effectively learn from the teacher’s supervision. On-policy KD’s low quality sample can be found in Sec 5.3 and Appendix G.

**Table 2: SKD outperforms all baseline KDs with SFTed student initialization of instruction following model and domain-specific fine-tuned model across three text generation tasks. We use supervised SFTed Gemma-7B-it and Qwen2-7B-it model as teachers and supervised SFT Gemma-2B-it and Qwen2-0.5B-it as students. SKD is performed under $K$=$25$.**
|  | KD initiates with SFTed Gemma-2B | KD initiates with SFTed Qwen-0.5B |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | Translation | Summarization | GSM8K | Translation | Summarization | GSM8K |
| Baselines | COMET | ROUGE-L | Acc | COMET | ROUGE-L | Acc |
| SFT | 72.5 | 31.7 | 18.7 | 57.6 | 29.2 | 31.8 |
| Supervised KD | 73.5 | 34.4 | 22.0 | 59.3 | 32.3 | 31.9 |
| On-policy KD | 74.6 | 34.8 | 28.8 | 58.3 | 32.1 | 32.2 |
| ImitKD | 74.6 | 34.4 | 26.2 | 59.9 | 32.2 | 31.8 |
| SKD | 75.0 | 35.4 | 29.2 | 60.6 | 33.1 | 33.4 |

Figure: Figure 4: Comparison between SKD and baseline KD methods with different model initializations. We find that on-policy methods struggles when the student model starts with a poor initialization, leading to performance degradation and becoming stuck at a low level throughout training. On-policy KD requires student model to have a good initialization. In contrast, SKD outperforms supervised and on-policy KD under both IT and supervised FT Gemma-2B initialization. On-policy KD’s low quality sample can be found in Appendix G. SKD is performed under $K$=$25$.
Refer to caption: x5.png

As shown in Table 2, SKD with SFT-initialized students consistently outperforms all baselines across both model families. While on-policy KD can narrow the gap with SKD through a two-stage training process (supervised fine-tuning + KD) for Gemma-2B students, it lags behind SKD for smaller models like Qwen-0.5B. Additionally, in Section 5.4 and Appendix Section H, we demonstrate that SFT can lead to over-fitting in student models under low-data regimes, resulting in sub-optimal performance in post-training KDs. Furthermore, despite some advantages over supervised and on-policy KD with IT-initialized student models, ImitKD does not exhibit clear benefits with higher quality initialized student models. This indicates that SKD’s token-level mixing of teacher and student model tokens can achieve more optimal performance than sequence-level mixing.

**Table 3: SKD outperforms other KD approaches with IT and SFTed student initialization across three text generation tasks at 100 data points. We use supervised SFT Gemma-7B-it as teacher and instruction tuned Gemma-2B and supervised FT Gemma-2B-it as students. SKD is performed under $K$=$25$.**
|  | KD initiates with IT Gemma-2B | KD initiates with SFTed Gemma-2B |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | Translation | Summarization | GSM8K | Translation | Summarization | GSM8K |
| Baselines | COMET | ROUGE-L | Acc | COMET | ROUGE-L | Acc |
| SFT | 59.8 | 23.0 | 8.11 | 59.8 | 23.0 | 8.11 |
| Supervised KD | 65.8 | 31.7 | 17.0 | 60.9 | 32.0 | 14.4 |
| On-policy KD | 33.7 | 32.2 | 12.1 | 63.4 | 32.4 | 16.1 |
| SKD | 67.0 | 32.3 | 18.8 | 63.8 | 32.6 | 16.7 |

### 5.4 Low data regime

We further investigated SKD under an extremely low data setting with 100 data points per task (10% of the training data). As shown in Table 3, SKD outperformed supervised KD and on-policy KD for all three tasks when initialized with IT and SFTed Gemma-2B, demonstrating SKD’s superior performance in this extreme low-data setting. Notably, both supervised KD and SKD exhibited lower or comparable performance when switching from instruction-tuned to SFTed initialization. This is attributed to unavoidable overfitting during SFT with only 100 data points. In Appendix Section H, we illustrate how overfitting at the SFT stage can lead to suboptimal performance in post-training KDs. This finding suggests that end-to-end SKD, capable of training from any base model and bypassing the SFT stage, offers a more optimal solution than two-stage on-policy KD (SFT+KD) in extreme low-data scenarios.

Figure: Figure 5: Ablation study. We conduct an ablation study by training the student model with supervised KD for the first half and self-generated samples for the second half. Our results demonstrate that this mixed training strategy can mostly be outperformed by either supervised KD or on-policy KD across three tasks and under two model initialization. Moreover, it is significantly outperformed by our SKD. We excluded on-policy KD from the leftmost figure due to its extremely low COMET score.
Refer to caption: x6.png

### 5.5 Two-staged training

Given that SKD transitions from supervised KD-like behavior at the beginning of training to on-policy KD towards the end, we investigated whether a simple approach of training with supervised KD for the first half of iterations and self-generated samples for the second half could achieve SKD-like performance. We conducted an ablation study by training the student model with supervised KD first, followed by on-policy KD (see Appendix Section I for details). As shown in Figure 5, SKD significantly outperforms this mixed training strategy. Surprisingly, this heuristic baseline was often outperformed by either supervised KD or on-policy KD across three tasks and two model initializations. This is likely due to the data discrepancy encountered during the training process – naively mixing samples from the teacher and student in two different stages lacks natural confidence measurement of the tokens; therefore, simply combining supervised and on-policy KD can even degrade the performance of either approach. This underscores the importance of our adaptive SKD, which implicitly covers intermediate transitions between two model behaviors.

### 5.6 Summary and Additional results

SKD consistently outperforms baselines in both task-specific and task-agnostic distillation. Unlike on-policy methods, which are sensitive to student-generated sample quality, SKD maintains its advantage even with low-quality student model initialization. As shown in Appendix H, overfitting in SFT can hinder post-training KD’s performance. SKD’s end-to-end approach, bypassing SFT, provides a significant benefit, especially in low-data scenarios. Merging supervised and on-policy KD naively can lead to performance degradation due to data inconsistencies.

Speculative decoding, as described in (Leviathan et al., ), is a crucial technique for accelerating LLM inference. The quality of the draft model used in this process directly influences the speed-up achieved. As demonstrated in Appendix J, our SKD-trained draft model significantly outperforms the instruction-tuned base model in terms of token acceptance rate ($71$% to $85$% higher), leading to a $1.2$X speed-up in speculative decoding.

## 6 Related Work

Knowledge Distillation (KD) (Hinton et al., ) is a widely adopted technique to transfer knowledge from a large, complex teacher model to a smaller, more efficient student model. This approach has been successfully applied to autoregressive models (Sanh et al., ) and language models in particular (Yang et al., ). There are two primary categories of KD methods: hard label and soft label. Hard label methods train the student model to mimic the teacher’s input-output pairs, similar to supervised fine-tuning (Kim & Rush, ; Taori et al., ; Chiang et al., ; Xu et al., ). Various data augmentation techniques, such as fine-tuning on rationales (Hsieh et al., ) or instruction-following formats (Wang et al., ), have been explored within this framework. Soft label methods, on the other hand, compute logits for both the student and teacher models and minimize a distance metric, such as KL-divergence, between their output distributions (Hinton et al., ; Sanh et al., ). Samples for computing logits can be generated by either model or both. Our proposed SKD falls under the category of soft label KD. To the best of our knowledge, this work is the first to introduce interleaved teacher-student sampling for soft-label KD.

Soft-label Knowledge Distillation can be broadly categorized into supervised KD and on-policy KD based on the source of training data. In supervised KD (Sanh et al., ), the training data is fixed, leading to a discrepancy between training and inference time for the student model (Agarwal et al., ). MiniLLM (Gu et al., ), for instance, frames this as a reinforcement learning problem, optimizing reverse KL divergence over the fixed dataset. Lin et al. (). explored different distance metrics, including JSD and total variation distance. On the other hand, on-policy KD, pioneered by ImitKD (Lin et al., ), leverages samples generated by the student model itself. This approach addresses the training-inference mismatch by directly sampling target tokens from the student’s own output. Inspired by on-policy imitating learning (Ross et al., ), Agarwal et al. () introduced on-policy KD for language models, calculating both teacher and student token probabilities to train the student on its self-generated data. DistillM (Ko et al., ) further enhanced this by incorporating a mix of on-policy samples and ground truth. As discussed in §3.2, both supervised KD and on-policy KD are special cases of SKD. SKD outperforms both baselines and surpasses the performance from mixing ground-truth and on-policy samples.

Speculative Decoding (Leviathan et al., ; Chen et al., ) inspired our interleaved sampling approach in SKD. However, instead of sampling from teacher’s distribution, SKD assesses the feasibility of student-proposed tokens within the teacher’s top $K$ tokens. Interleaved sampling allows us to obtain high quality on-policy samples. Zhou et al. (); Liu et al. () enhanced speculative decoding alignment with various KD methods. Building on these, we benchmarked SKD’s speedup against KD baselines, finding it superior in both speed and token acceptance.

## 7 Conclusion

We propose Speculative Knowledge Distillation (SKD), a novel method that addresses the limitations of existing KD approaches. SKD leverages interleaved sampling to mitigate the discrepancy between training and inference samples and remove low-quality student-generated data. Our experiments consistently show SKD’s superiority in both task-specific and task-agnostic distillations. SKD is robust to various model families, initialization, and dataset sizes. We also demonstrate its superiority in token-level mixing over sequence-level mixing. Furthermore, SKD-trained student models can accelerate speculative decoding, enabling new applications.

## 8 REPRODUCIBILITY STATEMENT

We discussed our model implementation, hyper-parameters, baseline details in Section 4. We include information about our dataset and evaluation metrics in Section 4.1. We discuss training and hyper-parameter details in Appendix C. We include examples of our prompt and response for all tasks in Appendix E. We include additional information about our evaluation metrics in Appendix F. All data, model and evaluation metrics are open sourced and available for research community. Upon the camera ready of the paper, we will include a link to a publicly accessible code repository, including all implementations about baseline KDs and SKD, running instructions and hyper-parameter settings. All results reported in this paper can be fully reproduced by the community.

## Appendix A Divergence Metrics.

As previously discussed, for each token position, the auto-regressive language model (LM) $M$ can generate a probability distribution, $p(.|X,y_{<k})$. Assuming the teacher and student models share the same vocabulary $V$, we can directly employ divergence metrics, like KL divergence, to quantify the distance between their respective probability distributions. For token position $k$, let $P(C)$ and $Q(C)$ represent the vocabulary distributions of the teacher and student, respectively. The KL divergence at this position is calculated as $D_{KL}(P||Q)=\sum_{c\in C}P(c)(\log\frac{P(c)}{Q(c)})$. By using this divergence metric as our objective function, we can effectively conduct knowledge distillation between the teacher and student models.

## Appendix B Can SKD outperform baselines under different acceptance criteria?

In this section, we tested a wide range of $K$ values and showed SKD’s performance across three text generation tasks. In Figure 6, we show that while extremely low or high $K$ values can not yield the most optimal performance, a broad range of $K$ values between $5$ to $50$ can consistently outperform supervised KD and on-policy KD across all three text generation tasks. The overall performance trend confirms our hypothesis that as $K$ increases, SKD degenerates to on-policy KD like behavior, while decreasing $K$ degenerates SKD to supervised KD like behavior. Consequently, neither the highest nor lowest values of $K$ are optimal.

We maintained a fixed $K$ of 25 (a common choice for top-K sampling) in our primary experiments to assess SKD’s general effectiveness without extensive parameter tuning. Our results in this section suggest that further task-specific optimization of $K$ could potentially yield even better performance.

Figure: Figure 6: We test a wide range of $K$ values and show SKD’s performance across three text generation tasks. We initialize student model with Gemma-2b-it. We use SKD with $K$=$25$ (a common choice for top-k sampling) for our experiments in prior sections. In this Figure, our results demonstrate that while extremely low or high $K$ values can not yield the most optimal performance, a broad range of $K$ values between $5$ and $50$ consistently outperforms supervised KD and on-policy KD across all three text generation tasks.
Refer to caption: x7.png

## Appendix C Training and Hyperparameter Details for Baselines and SKD

For both Gemma-2B-it and Qwen2-0.5B-it models, we set the learning rate to 1e-5. We disabled dropout during sampling, as in direct preference learning (Rafailov et al., ). Maximum input and output lengths were set task-specifically: (256, 256, 1024, 1024) and (256, 512, 128, 1024) for translation, arithmetic reasoning, summarization, and math instruction, respectively. All baseline models used a batch size of 8 and a gradient accumulation step of 1. Student model is trained over $375$, $375$, $375$ and $1225$ training steps for summarization, GSM, math and translation tasks respectively. We train $7500$ steps for full data training at Math. For extreme low data setting (with prompt size 100), we trained student model over $125$ steps for summarization, GSM and translation tasks. We report checkpoint results with the lowest validation loss.

Initial experiments with top-p sampling (temperature=$1$, top-p=$1$) led to significantly degraded generation quality compared to greedy decoding (temperature=0 0). This was particularly evident in math instruction task, where accuracy halved. To balance exploration and performance, we gradually reduced temperature and top-p until top-p sampling matched or outperformed greedy decoding. We ultimately settled on temperature=$0.5$ and top-p=$0.5$. This setting effectively limits the answer space while allowing for sufficient exploration. To ensure fair comparison, we kept the student temperature and top-p parameters identical for on-policy KD and SKD.

Two additional hyperparameters, teacher model top-p and temperature, can influence SKD’s performance. When a student token is rejected by the teacher, these parameters affect the resampled teacher token. Similar to prior work (Agarwal et al., ), we found that lower teacher temperatures improve distillation. We conducted a grid search over top-p $\{0.5,1\}$ and temperature $\{0.2,0.3,0.4\}$ for translation, arithmetic reasoning, and summarization on respective validation sets. While the optimal temperature varied slightly across tasks, $0.2$ generally performed well. Although a full vocabulary distribution (top-p=$1$) typically yielded the best results, reducing top-p to $0.5$ improved ROUGE-L from $0.5$ to $0.9$ across different temperatures at Summarization task. Therefore, we set temperature to $0.2$ for all tasks and top-p to $0.5$ for summarization and $1$ for other tasks. All hyperparameter settings are consistent between Gemma-2B and Qwen2-0.5B-it model.

In Section B, we explore a wide range of $K$ values and test SKD’s performance across three text generation tasks. Our results demonstrate that a wide range of $K$ values can outperform baseline knowledge distillation methods, such as On-policy KD and Supervised KD. To ensure consistent evaluation, we selected a constant $K$ value of $25$ (a common choice for Top-K sampling) for all tasks and ablation studies. However, future research could explore more dynamic $K$ selection strategies to potentially further enhance SKD performance.

**Table 4: In this Table, we show performance of Gemma-2B-it (student), supervised FT Gemma-2B-it (student) and supervised FT Gemma-7B-it (teacher) at both task specific distillation and task agnostic distillation.**
|  | Task specific distillation | Task agnostic distillation |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
|  | MT | Summ | GSM8K | Math | GSMplus | SVAMP | ASDiv |
| Baselines | COMET | ROUGE-L | Acc | Acc | Acc | Acc | Acc |
| Gemma-2B-it | 44.5 | 11.7 | 0 | 0.68 | 0.5 | 1.6 | 1.8 |
| Supervised FT Gemma-2B-it | 72.5 | 31.7 | 18.7 | 4.82 | 4.70 | 12.8 | 15.2 |
| Supervised FT Gemma-7B-it | 79.2 | 36.8 | 60.0 | 19.5 | 40.5 | 66.6 | 76.1 |

**Table 5: In this Table, we show performance of Qwen-0.5B-it (student), supervised FT Qwen-0.5B-it (student) and supervised FT Qwen-7B-it (teacher) at both task specific distillation and task agnostic distillation.**
|  | Task specific distillation | Task agnostic distillation |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
|  | MT | Summ | GSM8K | Math | GSMplus | SVAMP | ASDiv |
| Baselines | COMET | ROUGE-L | Acc | Acc | Acc | Acc | Acc |
| Qwen-0.5B-it | 44.9 | 15.2 | 0 | 7.0 | 3.9 | 1.7 | 3.9 |
| Supervised FT Qwen-0.5B-it | 57.6 | 29.2 | 31.8 | 2.6 | 3.9 | 10.5 | 14.1 |
| Supervised FT Qwen-7B-it | 79.4 | 38.1 | 73.8 | 46.0 | 60.4 | 84.8 | 90.0 |

**Table 6: This table presents baseline and SKD performance across four held-out test sets. The teacher model is a supervised FT Gemma-7B, while the student model is an instruction-tuned Gemma-2B. The left side of the table shows results for all baselines and SKD using 1,000 prompts from our constructed training set. The right side displays results using 10,000 prompts.**
|  | KD initiates with Gemma-2B-it (1k prompts) | KD initiates with Gemma-2B-it (10k prompts) |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Math | GSMplus | SVAMP | ASDiv | Math | GSMplus | SVAMP | ASDiv |
| Baselines | Acc | Acc | Acc | Acc | Acc | Acc | Acc | Acc |
| SFT | 4.82 | 4.70 | 12.8 | 15.2 | 10.9 | 16.8 | 42.3 | 56.8 |
| Supervised KD | 7.92 | 12.1 | 42.7 | 53.4 | 10.7 | 16.1 | 46.7 | 63.3 |
| On-policy KD | 7.92 | 15.1 | 44.5 | 57.4 | 10.3 | 18.6 | 49.9 | 61.8 |
| SKD | 9.56 | 16.9 | 46.9 | 59.1 | 10.2 | 19.7 | 50.4 | 68.4 |

## Appendix D Additional results

Tables 4 and 5 show the performance of Gemma-2B-it (student), Qwen-0.5B-it (students) and supervised FT teacher models (Gemma-7B-it and Qwen-7B-it) under task-specific and task-agnostic distillation. Table 6 presents task-agnostic instruction tuning results for baseline and SKD models across four held-out test sets, using both low-data (1k prompts) and full-data (10k prompts) training regimes. These results align with those shown in Figure 3.

## Appendix E Task input and output Format

We include input and output formats for four text generation tasks that we studied in this paper: math instruction following (Table 7), arithmetic reasoning (Table 8), dialogue summarization (Table 9) and translation (Table 10).

**Table 7: This outlines the input and output format for math instruction following task. The output includes chain-of-thought steps leading to the computed answer, which is presented in a boxed format. We evaluate these answers by comparing their numerical equivalence to the ground truth.**
| Prompt for model: |
| --- |
| Solve the following math problem step-by-step.<br>Simplify your answer as much as possible. Present your final answer as boxed {Your Answer}.<br>For a positive integer $n$, let<br>$H_{n}=1+\frac{1}{2}+\frac{1}{3}+\dots+\frac{1}{n}$ Compute<br>$\sum_{n=1}^{\infty}\frac{1}{(n+1)H_{n}H_{n+1}}$. |
| Output for model: |
| Step 1: First, let’s understand what the problem is asking for. We are given the expression $H_{n}=1+\frac{1}{2}+\frac{1}{3}+\dots+\frac{1}{n}$ for a positive integer n, and we need to compute the sum $\sum_{n=1}^{\infty}\frac{1}{(n+1)H_{n}H_{n+1}}$. |
| Step 2: To begin solving this problem, we can start by manipulating $\frac{1}{(n+1)H_{n}H_{n+1}}$ to see if we can simplify it in any way. By doing this, we can rewrite $\frac{1}{(n+1)H_{n}H_{n+1}}$ as $\frac{H_{n+1}-H_{n}}{H_{n}H_{n+1}}$. |
| Step 3: This allows us to express the given sum as a telescoping series, where each term of the sum cancels out most of the preceding and following terms. |
| Step 4: Taking the sum $\sum_{n=1}^{\infty}\frac{1}{(n+1)H_{n}H_{n+1}}$ in parts, we can break it down term by term to simplify the expression and find the pattern of the telescoping series. |
| Step 5: By expressing each term of the sum as a difference involving the harmonic numbers $H_{n}$, we can observe how the terms cancel out, and this leads us to the result of the sum. |
| Step 6: To validate and confirm the result, we need to acknowledge the fact that $H_{n}\to\infty$ as $n\to\infty$, and we can prove this by showing that the sum of the harmonic series diverges. |
| Step 7: This proof involves demonstrating that the sum of the harmonic series increases without bound as n increases, which verifies the behavior of $H_{n}$ as $n\to\infty$. |
| Step 8: Therefore, the final answer to the given problem is $\boxed{1}$, and it is backed by the property of harmonic numbers tending to infinity as the terms increase. |
| Answer: |
| $\boxed{1}$ |

**Table 8: This outlines the input and output format for arithmetic reasoning task. The output includes chain-of-thought steps leading to the computed answer, which is presented after ”####”. We evaluate these answers by comparing their numerical equivalence to the ground truth.**
| Prompt for model: |
| --- |
| Q: American carmakers produce 5 650 000 cars each year. Then they distribute it to 5 car suppliers. The first supplier receives 1 000 000 cars each. The second supplier receives 500 000 cars more while the third car supplier receives the same number as the first and second suppliers combined. The fourth and the fifth supplier receive an equal number of cars. How many cars do the fourth and fourth suppliers each receive? |
| Output for model: |
| The second supplier receives 1 000 000 + 500 000 = 1 500 000 cars. |
| The third supplier receives 1 000 000 + 1 500 000 = 2 500 000 cars. |
| So the first to third suppliers receive a total of 1 000 000 + 1 500 000 + 2 500 000 = 5 000 000 cars. |
| Thus, 5 650 000 - 5 000 000 = 650 000 cars are left to be divided by the two other suppliers. |
| Therefore, the fourth and fifth suppliers receive 650 000/2 = 325 000 cars each. |
| #### 325,000 |

**Table 9: This outlines the input and output format for dialogue summarization task. The output is the summary of above dialogue. We evaluate summary by computing its ROUGE-L score to the ground truth.**
| Prompt for model: |
| --- |
| #Person1#: Hi, Mr. Smith. I’m Doctor Hawkins. Why are you here today? |
| #Person2#: I found it would be a good idea to get a check-up. |
| #Person1#: Yes, well, you haven’t had one for 5 years. You should have one every year. |
| #Person2#: I know. I figure as long as there is nothing wrong, why go see the doctor? |
| #Person1#: Well, the best way to avoid serious illnesses is to find out about them early. So try to come at least once a year for your own good. |
| #Person2#: Ok. |
| #Person1#: Let me see here. Your eyes and ears look fine. Take a deep breath, please. Do you smoke, Mr. Smith? |
| #Person2#: Yes. |
| #Person1#: Smoking is the leading cause of lung cancer and heart disease, you know. You really should quit. |
| #Person2#: I’ve tried hundreds of times, but I just can’t seem to kick the habit. |
| #Person1#: Well, we have classes and some medications that might help. I’ll give you more information before you leave. |
| #Person2#: Ok, thanks doctor. |
| Output for model: |
| Mr. Smith’s getting a check-up, and Doctor Hawkins advises him to have one every year. Hawkins’ll give some information about their classes and medications to help Mr. Smith quit smoking. |

**Table 10: This outlines the input and output format for translation task. The output is the English translation. We evaluate translation by computing its COMET score to the reference and source.**
| Prompt for model: |
| --- |
| Translate Assamese sentence into English. This is the Assamese sentence: [Assamese Texts]. Here is the English translation: |
| Output for model: |
| It was to last for the next 40 years and would be fought for real, by proxy armies, on battlefields from Africa to Asia, in Afghanistan, Cuba and many other places. |

## Appendix F Evaluation Metrics

As in (Freitag et al., ), we adopted the learned evaluation metric COMET (Rei et al., ) for assessing our translation outputs, instead of the n-gram-based BLEU (Papineni et al., ). Specifically, we used the comet-da-22 checkpoint (Rei et al., ), a widely recognized metric that exhibits strong correlation with human judgments at WMT22 (Freitag et al., ). COMET is a learned metric trained on a large-scale dataset of human judgments for evaluating translation quality. It takes source, translation, and reference as input, producing a quality score between 0 0 and $1$. A COMET score improvement of $0.5$ to $1.0$ indicates a substantial improvement.

Tables 7 and 8 illustrate the structured output format used to derive final answers from math instruction tasks and GSM-8k. We employed the evaluation code from (Yuan et al., ) to assess the mathematical equivalence of parsed answers and ground truths. This involved converting mathematical symbols like $\pi$ into their numerical values for more accurate numerical comparisons, rather than relying solely on string matching. Implementation details can be found in the original paper (Yuan et al., ). We reported accuracy for both tasks.

We evaluated summarization output against references using ROUGE-L (Lin, ). ROUGE-L ranges from 0 to 100, with a 1-point increase indicating substantial improvement.

**Table 11: We sample training data from SKD and on-policy methods. Our results show that solely sampling from the student model can lead to hallucinated translations and limited training signals (On-policy output). Conversely, SKD generates fluent, semantically similar outputs while staying close to the student’s distribution (SKD output). Student model is supervised FT Qwen-0.5B.**
| Prompt for model: |
| --- |
| Translate Assamese sentence into English. This is the Assamese sentence: [Assamese Texts]. Here is the English translation: |
| Reference translation: |
| Italian is also the everyday language used by most of those who work in the state while Latin is often used in religious ceremonies. |
| On-policy student sampled output (COMET: 53): |
| The Estonian language is also used in the Estonian-speaking countries as a national language, and is used in the Estonian-speaking countries as a medium of instruction in the Estonian-speaking countries. |
| SKD student sampled output (COMET: 80): |
| The Latin language is still used in the country to a large extent, though it is not the official language. |

## Appendix G Case Study between On-policy and SKD samples

Tables 11 and 12 illustrate low-quality training examples generated by on-policy models. In Table 11, the student model hallucinates ”Estonian language,” a nonexistent term in the source, leading to a poor translation. SKD with teacher’s interleaved sampling avoids this issue. In Table 12, on-policy KD samples repeat the n-gram ”in the past,” producing nonsensical translations. SKD’s interleaved sampling helps prevent this pathology. In Table 13, we sampled training data from SKD and on-policy methods using the Gemma-2B-it model. These samples were collected during the early stages of training. The on-policy sample exhibited repetitive outputs and errors in arithmetic calculations. While our SKD model also produced incorrect answers, incorporating teacher’s corrections enabled it to perform some intermediate steps correctly and maintain accurate arithmetic calculations. In Table 14, we sampled training data from SKD and on-policy methods using the Gemma-2B-it model. These samples were collected during the late stages of training. We observed that teacher model can perform some intermediate correction on the reasoning step, ultimately leading to correct answer.

**Table 12: We sample training data from SKD and on-policy methods. Our results show that solely sampling from the student model can lead to poor translations and limited training signals (On-policy output). Conversely, SKD generates fluent, semantically similar outputs while staying close to the student’s distribution (SKD output). We omit source sentence due to font format issues. Student model is supervised FT Qwen-0.5B.**
| Prompt for model: |
| --- |
| Translate Assamese sentence into English. This is the Assamese sentence: [Assamese Texts]. Here is the English translation: |
| Reference translation: |
| It was to last for the next 40 years and would be fought for real, by proxy armies, on battlefields from Africa to Asia, in Afghanistan, Cuba and many other places. |
| On-policy student sampled output (COMET: 36): |
| This has been a long war, and has been fought in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in the past, in |
| SKD student sampled output (COMET: 72): |
| This has been going on for 40 years and is still ongoing, with the possibility of war, a civil war, Afghans in the Balkans, Ethiopia, Kenya, Somalia, and possibly even the United States. |

**Table 13: We sample training data from SKD and on-policy methods using the Gemma-2B-it model. These samples are collected during the early stages of training. The on-policy sample exhibit repetitive outputs and errors in arithmetic calculations. While our SKD model also produces incorrect answers, incorporating teacher’s corrections enables it to perform some intermediate steps correctly and maintain accurate arithmetic calculations.**
| Prompt for model: |
| --- |
| Q: Libby is building an igloo in her backyard using bricks of snow. She builds her igloo in rows, using a total of 10 rows of bricks of snow. The bottom half of the igloo has 12 bricks of snow in each row while the top half of the igloo has 8 bricks of snow in each row. How many bricks of snow did Libby use for her igloo? A:<br>Correct Answer: |
| In the bottom half of the igloo, there are 10 rows / 2 = 5 rows of bricks of snow. |
| So in the bottom half of the igloo, there are 5 rows * 12 bricks of snow in each row = 60 bricks of snow. |
| In the top half of the igloo, there are 5 rows * 8 bricks of snow in each row = 5*8=4040 bricks of snow. |
| There is therefore a total of 60 + 40 = 100 bricks of snow in the igloo. #### 100 |
| On-policy student sampled output: |
| $35-$160= $125 |
| Libby has $160-$125=$35 |
| Libby has 160-35= $125 |
| Libby has 160-35= $125… |
| SKD student sampled output: |
| The bottom half of the igloo has 12 x 10 = 120 bricks of snow. |
| The top half of the igloo has 8 x 10 = 80 bricks of snow. |
| Libby used 120 + 80=200 bricks of snow. #### 200 |

**Table 14: We sample training data from SKD and on-policy methods using the Gemma-2B-it model. These samples are collected during the late stages of training. We observe that teacher model can perform some intermediate correction on the reasoning step, ultimately leading to correct answer.**
| Prompt for model: |
| --- |
| Q: John Smith buys 3 cakes for $12 each and splits the cost with his brother. How much did he pay? A:<br>Correct Answer: |
| He spent 3*12=36 total<br>So he paid 36/2=18<br>#### 18 |
| On-policy student sampled output: |
| Each brother pays $12/3 = 4 each |
| So he pays $12 - $4 = 8 #### 8 |
| SKD student sampled output: |
| The 3 cakes cost John 3 x $12 = $36 |
| So he paid $36 / 2 = $18 each #### 18 |

## Appendix H Over-fitting at SFT stage can cause post-training KDs to sub-optimal performance

Figure 7 demonstrates that overfitting during the SFT stage can hinder optimal performance in post-training KD. End-to-end KD methods, such as SKD, which can train from arbitrary base model and bypass the SFT stage, offer significant benefits. To investigate this, we selected three checkpoints from the supervised FT stage of the student model, corresponding to training steps 64, 128, and 192. Based on validation, the checkpoint at step 128 represents the converged model, while step 192 exhibits signs of overfitting (Figure a). Figure (b) shows that the converged checkpoint (step 128) outperforms the overfitted checkpoint (step 192). Moreover, we found that training from an earlier checkpoint (step 64) can yield better results than the overfitted checkpoint (step 192). This observation is consistent for both on-policy and SKD. However, SKD is an end-to-end training pipeline. Our findings confirm that overfitting during the SFT stage can lead to suboptimal performance in post-training KD. As shown in Section 5.4, training with extremely limited data can inevitably lead to overfitting during the SFT stage, further highlighting the superiority of our end-to-end SKD pipeline compared to two-stage KD approaches like on-policy KD.

Figure: Figure 7: In this Figure, we demonstrate that over-fitting during the SFT stage can hinder optimal performance in post-training knowledge distillation. End-to-end KD methods, such as SKD, which can train directly from arbitrary base model and bypass the SFT stage, offer significant benefits. Figure (a) presents three checkpoints from the SFT stage, corresponding to training steps $64$, $128$, and $192$. Based on validation loss, the checkpoint at step $128$ represents the converged model, while step $192$ shows signs of overfitting. Figure (b) demonstrates that the converged checkpoint (step $128$) outperforms the overfitted checkpoint (step $192$). This finding is consistent for both on-policy KD and SKD.
Refer to caption: x8.png

## Appendix I Training details of mixed training of supervised KD and on-policy KD

The parameter setting for this baseline is identical for the setup in Appendix C. We start to train student model with supervised KD samples for $188$ steps (roughly $1.5$ epoches), then we perform on-policy KD for another $187$ steps. Total training step is $375$. Initial learning rate is $1e$-$5$ and linear decay is applied during training process.

## Appendix J speedup to speculative decoding

Motivated by (Zhou et al., ), we further investigate the speedup achieved by each baseline model when used as a draft model to approximate a teacher model during speculative decoding (Leviathan et al., )]. As depicted in Figure 8, SKD-trained student models consistently exhibit the highest speedup and token acceptance ratio (Approximation model will accept or reject tokens depending on whether draft token is sampled under its probability distribution) across translation and summarization tasks. These models accelerate instruction-tuned Gemma-2B by 1.21x and 1.17x in translation and summarization, respectively, while also improving token acceptance ratio by 1.71x and 1.85x. These results demonstrate that SKD not only produces superior performance but also generates the closest distribution to the teacher model. This opens up a new application: utilizing SKD to train powerful draft models and accelerate the speculative decoding process.

Figure: Figure 8: We utilize student models trained with various baseline KD methods as draft models and a supervised fine-tuned Gemma-7B-it model as the approximation model for speculative decoding. Our results demonstrate that SKD-trained student models significantly outperform the others, achieving the highest speedup and token acceptance ratio. SKD accelerates Gemma-2B-it by 1.21x and 1.17x in translation and summarization, respectively, while also improving token acceptance ratio by 1.71x and 1.85x.
Refer to caption: x10.png

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 6
presentation: 4
questions: - Is there a way to show how severe the OOD issue of on-policy KD is? For example, would it be possible to quantitatively evaluate?
- Why does the setting proposed in Section 5.2 a "task-agnostic" one?
soundness: 3
strengths: - The motivation of the paper is convincing. Supervised KD is known to suffer from distribution shift, while on-polich KD suffers from the OOD problem.
- The organization is well and the paper is carefully written. The problem setting is well formulated.
- All-around empirical results pretty much validates the effectiveness of the method. The authors conducted experiments on four tasks and one task agnostic distillation.
summary: This paper proposed a new knowledge distillation method called Speculative KD (SKD), inspired by speculated decoding, that addresses the drawbacks from supervised KD and on-policy, while keeping the advantages of both worlds. Specifically, to circumvent the issue of supervised KD, SKD enforces the teacher model to evaluate the output of the student model. To alleviate the "cold-start" OOD problem of on-policy KD, SKD filters out intermediate tokens that the teacher is unlikely to generate, and re-samples from the teacher.
weaknesses: - Although it is relatively straightforward to agree with the authors that the on-policy KD suffer from the OOD issue, is there a way to show how severe the issue is? For example, would it be possible to quantitatively evaluate?
- It is a bit hard to understand the task-agnostic distillation setting. I see the training and test sets are both based on the math reasoning datasets. Where does the shifts come from?
- In table 2, it is obvious that the improvements are rather marginal. In Table 1, it seems the proposed method falls short behind another baseline. Would you explain?
- Significant tests for Table 1 & 2 could further strengthen the results.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 4
questions: Question 1: The paper does not analyze how early low-quality samples generated by the student model may impact complex tasks, such as mathematical reasoning. How do the authors believe these early samples could affect the stability of teacher feedback and convergence rates?

Question 2: Although the authors' framework incorporates dynamic adjustments and adaptive switching between teacher and student sampling, how do they envision customizing these mechanisms for different task types or complexities?

Question 3: While the authors examine SKD’s performance under various initialization conditions, how does initial model quality affect convergence efficiency and stability?
soundness: 3
strengths: 1. The paper introduces Speculative Knowledge Distillation (SKD) to address the distribution mismatch between training and inference in traditional knowledge distillation. By combining on-policy sampling with speculative decoding, SKD enables efficient knowledge transfer, allowing the student model to autonomously generate and refine its samples. 

2. The authors validate SKD across various task scenarios, including translation, summarization, arithmetic reasoning, and instruction following, covering both task-specific and task-agnostic applications. Through comparisons with multiple baseline methods (e.g., supervised KD, on-policy KD, and ImitKD), the experimental setup is comprehensive and broadly applicable, providing certain support for SKD’s performance.

3. SKD’s design builds on imitation learning theory, lending theoretical validity to the approach. Additionally, the paper includes detailed implementation steps and hyperparameter settings, providing actionable guidance for future research.
summary: The authors present Speculative Knowledge Distillation (SKD), a new KD method for effective knowledge transfer from large teacher models to smaller student models. SKD combines on-policy sampling with speculative decoding, allowing students to propose tokens while the teacher verifies and refines them, addressing issues in traditional KD methods like distribution mismatches and low-quality samples.

The authors evaluate SKD across tasks including translation, summarization, math, and instruction following, consistently outperforming baseline methods across different domains, data sizes, and initialization strategies. SKD shows robust performance in both task-specific and task-agnostic distillation, with strong generalization on complex math tasks and advantages in low-data scenarios.

Contributions:
1. Introducing SKD, which integrates on-policy sampling and speculative decoding for adaptive, high-quality knowledge transfer.
2. Demonstrating SKD’s superior performance across varied tasks, outperforming supervised and on-policy KD.
3. Validating SKD’s adaptability and efficiency, especially in diverse model initializations and low-data settings, without requiring additional supervised fine-tuning.
weaknesses: 1. Impact of Early Low-Quality Samples Lacks In-Depth Analysis: The SKD framework uses interleaved sampling and speculative decoding to handle low-quality samples generated by the student model during its initial training stages. However, the authors do not analyze how these low-quality samples may impact complex tasks, such as mathematical reasoning. Early low-quality samples could affect the stability of teacher feedback, slow convergence, and compromise overall stability. The authors could consider adding an analysis of model convergence behavior across tasks of varying complexity in their experiments to gain a more comprehensive understanding of SKD’s convergence characteristics in complex scenarios.

2. Task-Specific Adaptability: SKD maintains overall sampling quality through dynamic adjustments of token-level sample quality and adaptive switching between teacher and student sampling. However, the authors do not explore how to customize these mechanisms for different task types or complexities.

3. Effect of Model Initialization Lacks Discussion: While the authors examine SKD’s performance under different initialization conditions, such as instruction tuning and supervised fine-tuning, they do not discuss how initial model quality affects SKD’s convergence efficiency and stability. The authors might consider investigating convergence behavior under various initialization qualities, particularly in low-quality scenarios, by dynamically adjusting the proportion of teacher samples to stabilize the training process.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 4
questions: **(a)** It is not 100% clear to me why would one replace a `bad token` with a token sampled from teacher's distribution. Could the authors give some intuition of why not re-attempt the sampling from student's distribution (line 4 in Algorithm 1) to replace the last token but under the condition that it's within TopK of the teacher? In otherwords, maybe the `bad` token sampled from student's distribution can be replaced with a `better` token that is sampled from the same distribution, rather than one that comes from the teacher's. If it is within the computational budget, an experiment to demonstrate this could also help prove the point.


**(b)** If I understood correctly, your method does not add any computational overhead to the distillation process, as it is a simple probability check and a token replacement. If so, please assert this both in rebuttal and also in the paper. Otherwise, a clear discussion on computational overhead is missing.
soundness: 4
strengths: **(a)** I think the paper is extremely well written. Especially for someone who might not be fully familiar with how Knowledge Distillation is applied for the LLMs, I really appreciate how clearly the authors explain prior work and how they relate to each other. With only the first pass, I was able to grasp the ideas across related work and the main contribution of the work.


**(b)** I think the simplicity of method is one of its advantages. Simply checking the student token's quality, based on whether it's within TopK of the teacher's seems simple and intuitive. (see also Question a).


**(c)** The quantitative results, especially compared to On Policy KD, seem promising (please see the disclaimer in Weaknesses section).
summary: The method seems very simple and, to some extent, sounds like an interpolation between two distillation strategies. Specifically, in Supervised KD the next-token distribution of student and teacher are trained to be similar on a fixed text, while in On-policy it is done on student’s generated text. The method thus proposes to evaluate the tokens of the student and see if it’s within TopK of the teacher’s, if not, the generated tokens of the student is discarded and the rest is followed with teacher’s tokens. But if the generated token satisfies the TopK check, we continue with student’s tokens.


The idea is clearly motivated in the paper and makes sense. The quantitative results also seem promising (across tasks the method is mostly outperforming prior work by a good margin). Especially the method is consistently outperfoming On Policy KD, which supports the proposed idea of switching to teacher tokens if the student tokens are bad.
weaknesses: **(a)** I like the idea of evaluating the student’s suggestions by the teacher and selecting the Top-K based on the teacher’s probabilities. However, what concerns me is the computational cost of the method. Given that this is the main contribution of the work, I think it has to be more clearly studied. Specifically, what I really like to see is how often the line 5 in Algorithm 1 is triggered. In Lines 212-216 it is essentially claimed that this trigger rate is decreased over training time, I think it has to be evidenced with a figure (x axis being the training steps and y axis being the trigger rate). In a similar sense, I think it should be reported what is the final rate (i.e how likely it is for the student’s tokens to be in TopK of the teacher’s after the distillation).


**(b)** Moreover, I think the value for K should be studied. Essentially for high Ks the method should get somewhat similar to On Policy KD and for lower ones should be closer to Supervised KD? So there should be a sweet spot. Also I’m wondering if the value for K needs to be adjusted based on the training step? In other words, maybe higher or lower Ks would be better at the early or later stages of distillation? Some ablation experiments (with different numbers for K) could be helpful in addressing this.

------ Minor issues ------

**(c)** In Line 119, it is stated that “the temperature t `introduces` randomness”. I think this is not entirely correct as it’s the sampling that introduces randomness. The temperature adjusts the distribution so it `adjusts` the randomness (as it is also stated correctly in the following sentence).
 
**(d)** In Line 184, “Correct previous mistakes”,  it is not clear what `mistake` is referring to. Is a mistake a poorly generated Token? What does `correcting` refer to?

**(e)** Line 250: “Since our assumption `is` that”

**(f)** In the pseudo-code (Algorithm 1), some variables, such as alpha in line 3, are not defined.

**(g)** This is not really an issue, rather a suggestion. A tiny figure to explain concepts defined in Section 3.1 would be very helpful. Basically a 1-row figure divided into maybe 4-5 sub-columns, with every column showing how each of the prior work applies the loss at token-level. Maybe something like Figure 1 in https://arxiv.org/pdf/2104.09044 I think Figure (2) can be refactored a bit so that it also includes other prior work. 




------ Disclaimer ------

I’m not very familiar with KD approaches for LLM, and therefore it is possible that I haven't noticed a missing baseline. I would really appreciate it if other reviewers could also verify if the set baselines is complete.

Meta-review: The paper presents a Knowledge Distillation method that circumvents issues in both supervised KD and on-policy KD [Reviewer VY1t]

The authors evaluate SKD across several types of tasks including translation, summarization, math, and instruction following, consistently outperforming baseline methods across different domains, data sizes, and initialization strategies. SKD shows robust performance in both task-specific and task-agnostic distillation, with strong generalization on complex math tasks and advantages in low-data scenarios. [Reviewer cyWE]

The idea is clearly motivated in the paper and makes sense. The quantitative results also seem promising (across tasks the method is mostly outperforming prior work by a good margin). Especially the method is consistently outperfoming On Policy KD, which supports the proposed idea of switching to teacher tokens if the student tokens are bad. [Reviewer EjjZ]

Authors have addressed concerns from reviewers and have updated the manuscript with sufficient empirical results to substantiate claims and to enhance the robustness, scalability and convergence of the proposed SKD.

Decision: Accept (Poster)