Multi-Granularity Semantic Revision for Large Language Model Distillation

As illustrated by Eq. (1), prevalent KD methods agarwal2024policy ; gu2023minillm ; ko2024distillm , utilizes conditional sequences generated from the student model (denoted as $y\sim q_{\theta}(\cdot|x)$ ) for the distillation process. While these methods are designed to mitigate the training-inference mismatch between the fixed training data and the student’s auto-regressive inferences, they simultaneously risk introducing generation errors. Due to the limited capabilities of the student model, the generated sequences may contain additional errors which reduces the effectiveness of KD. To address this issue, we propose a sequence correction and re-generation (SCRG) strategy (shown in Fig. 2) to detect generation errors and re-generate sequences that steer the student model towards generating reliable and diverse sequences.

We denote the generated $n$-token sequence from the student model $q_{\theta}$ as $y^{s}_{<n+1}=(y^{s}_{1},y^{s}_{2},....,y^{s}_{n})$ which correspondences the probability outputs $(\hat{y}^{s}_{1},\hat{y}^{s}_{2},....,\hat{y}^{s}_{n})$, where $y^{s}_{i}\sim q_{\theta}(.|y^{s}_{<i},x)(1\leq i\leq n)$. Similarly, we denote the teacher model’s output sequence as $y^{t}_{<n+1}=(y^{t}_{1},y^{t}_{2},....,y^{t}_{n})$ and probability outputs $(\hat{y}^{t}_{1},\hat{y}^{t}_{2},....,\hat{y}^{t}_{n})$. We denote each token of the teacher model’s output sequence as $y^{t}_{i}\sim p(.|y^{s}_{<i},x)$. We follow previous methods agarwal2024policy ; gu2023minillm ; ko2024distillm using the student-generated outputs as the distillation dataset, and calculate token-wise KLD loss to evaluate the semantic cognitive differences between the teacher and student for each token to detect the position of the error token within the sequence $y^{t}_{<n+1}$. We formulate the detection process of the error token $y_{j}^{s}$ as

We then replace the $y_{j}^{s}$ by $y_{j}^{t}$ to construct new samples and re-generate the student output sequence and each token in $y^{s}_{<n+1}$ is formulated as

Our SCRG strategy can be seamlessly integrated with existing on-policy sampling agarwal2024policy and off-policy sampling ko2024distillm . By incorporating an adaptive scheduler ko2024distillm for student-model generation, we enhance the efficiency of our sampling process.

The probability output of LLMs is a high-dimensional vector for each token. However, existing modified Kullback-Leibler divergence (KLD) loss functions, used as knowledge distillation objectives, struggle to effectively capture the valuable distribution with high semantic knowledge from the teacher network. They either underfitt the the teacher’s distribution, as seen in forward KLD, or tend to overfit to a part of the high-probability region, as seen in reverse KLD. To address this issue, we design a Distribution-Adaptive Clipping Kullback-Leibler (DAC-KL) loss function (in Fig. 3) to capture high-density semantic regions of the teacher’s output probability distribution, which can be more easily imitated by the student models with limited capacity.

The probability outputs at the $i^{th}$ token position of both the teacher and student models are high-dimensional probability vectors with $M$ tokens, which are denoted as

We input these two probability vectors to a learnable MLP sub-network $f_{sub}$ to predict the upper limit quantile $u\in[0,1]$ and the lower limit quantile $l\in[0,u]$ of the probability distribution $\hat{y}^{t}_{i}$. We formulate this process as

where $\sigma(\cdot)$ is the SIGMOID activation, $sort(\cdot)$ is the decending sort operation, and $\mid$ represents the concatenation operation, $l$ is clipped into the range $[0,u]$.

The predicted quantiles $u$ and $l$ are used to adaptively clip out the high-density semantic classes from the teacher’s probability vector $\hat{y}^{t}_{i}$. We utilize the clipped high-density classes and the target class with the most probability value to construct a new probability vector $\hat{y}^{t*}_{i}$, which is formulated as

The high-density classes and the target class contain the most knowledge in the teacher’s probability distribution. Based on the corresponding positions of the clipped classes and target class of $\hat{y}^{t*}_{i}$, we construct the student’s new probability vector $\hat{y}^{s*}_{i}$. Then, we adopt a vanilla KLD to calculate the sum of token-wise distillation loss and the final loss is calculated on the dataset ($X$,$Y$):

where $L_{y^{s*}}$ is the length of the sequence generated from the proposed SCRG strategy.

Motivated by the work liu2022multi , we utilize the pre-defined chunker kiss2006unsupervised to extract spans (including noun phrases, verb phrases, and prepositional phrases) that have complete meanings from the input sequences, which split a sequence into several token sets. For each token in the input sequence, LLMs predict a high-dimensional probability vector. The relations between tokens within the same span should maintain consistent relations in the transformed probability space. Constraining the relation consistency between the outputs of the student and the teacher models is crucial to transfer semantic knowledge, as shown in Fig. 4.

We divide a probability sequence $\left[\hat{y}_{1},\hat{y}_{2},...,\hat{y}_{n}\right]$ into $n_{s}$ spans $s=\left[s_{1},s_{2},...,s_{n_{s}}\right]$ accoreding to the pre-defined span priors from $\left[y_{1},y_{2},...,y_{n}\right]$. Here, $s_{i}=\left[\hat{y}_{j},\hat{y}_{j+1},...,\hat{y}_{j+n_{s_{i}}-1}\right]$ represents $i^{th}span$, whcih starts at the $j^{th}$ token of the sequence and contains $n_{s_{i}}$ tokens. Both the student and teacher model outputs adhere to the same span priors for token divisions. Consequently, we divide the probability outputs of the student and teacher models into spans, denoting the $i^{th}$ span as

Next, we calculate the correlation between two adjacent tokens within the same spans and ensure consistency of this correlation between the probability outputs of the student model and the teacher model. To achieve this, we utilize the L2 distance to align the consistency. The span consistency loss is defined as follows:

where $\left|\cdot\right|_{2}$ represents the L2 distance function, and $\circ$ denotes the Hadamard multiplication operation calculating correlation in the high-dimensional probability space. It is important to note that the output sequence is also generated by the student using the SCRG strategy. For simplicity, we adopt a standard symbol representation for $\hat{y}^{t}_{j}$ and $\hat{y}^{s}_{j}$ instead of $\hat{y}^{t*}_{j}$ and $\hat{y}^{s*}_{j}$.

We use the proposed KD method in the SFT stage of based models. The student model is supervised by the distillation loss, guided by the finetuned teacher model, and also supervised by the SFT loss. The overall optimization objective for the student model is formulated as

where $\mathcal{L}_{\text{SFT}}$ represents the SFT loss, $\mathcal{L}_{\text{DAC-KLD}}$ represents the distillation loss using the DAC-KLD object, and $\mathcal{L}_{\text{span}}$ represents the span consistency loss which assists the distillation process.

We use the ROUGE-L lin2004rouge metric to evaluate the model-generated results and report the average scores of 5 generations for each prompt with different random seeds ($10$, $20$, $30$, $40$, $50$) for all test datasets. ROUGE-L evaluates the precision of the model’s output by measuring the longest common subsequence between the generated text and the reference text. It is well-suited for large-scale instruction-following evaluation due to its ability to capture both sentence-level structure and content.

We benchmark our method against several advanced knowledge distillation methods: (1) SFT Fine-tunes the student model on a fixed dataset in a vanilla manner; (2) KD  hinton2015distilling utilizes KLD on a fixed dataset; (3) SeqKD kim2016sequence fine-tunes on a teacher-generated dataset; (4) ImitKD lin2020autoregressive utilizes KLD on a dataset generated by the student model; (5) GKD agarwal2024policy utilizes Jensen-Shannon Divergence (JSD) agarwal2024policy on a mixture of a student-generated dataset and a fixed dataset; (6) MiniLLM gu2023minillm utilizes a policy gradient approach on a dataset generated by the student model; (7) DistiLLM ko2024distillm utilizes Skew KLD on a student-generated dataset sampling with an off-policy scheduler.

All of our baseline experiments are re-implemented using the open-source code ²²2https://github.com/jongwooko/distillm on the same GPU servers utilized by our method. Additionally, we execute these experiments using the exact hyper-parameters as specified in the original codebase.

For training the teacher and student models, we utilize four A100 (40GB) GPUs for the OpenLLAMA2, OPT, and GPT2 models and four A800 (80GB) GPUs for the LLAMA2 models. A fixed learning rate of 5e-4 is applied consistently across all experiments. Specifically, for the LLAMA2, OpenLLAMA2, and OPT models, we follow DistiLLM ko2024distillm , employing low-rank adaptation (LoRA) for the query and value weights with a rank of 16 for 10 epochs. In contrast, for the GPT2 models, we fine-tune all parameters for 20 epochs.

We present the quantitative comparison of state-of-the-art knowledge distillation methods evaluated using the ROUGE-L metric in Table 1. It is observed that:

(1) Our method outperforms existing methods in most distillation tasks, with only a few achieving second-best results, across five test datasets, including the LLAMA2, OPT, OpenLLAMA2, and GPT2 series of large language models. Particularly for the OPT datasets, our method shows an average score improvement of over 12% compared to the second-best performing methods.

(2) The KD methods, such as GKD, MiniLLM, and DistiLLM, utilizing student-generated datasets show a greater improvement in enhancing student performance compared to those using the fixed dataset. Furthermore, the distilled student models generally outperform the teacher models, which can be attributed to the mismatch between teacher-forcing training and free-run generation, i.e., exposure bias bengio2015scheduled . Our method can improve the performance of all student models on average scores of the five test datasets by at least 15%.

(3) We also provide some representative instruction-following cases in Section A.3, further highlighting the effectiveness and superiority of our method in achieving high-quality answers.

We conduct an ablation analysis of the proposed methods on the Dolly Validation set, Dolly Evaluation set and Self-Instruct dataset.