DA-KD: Difficulty-Aware Knowledge Distillation for Efficient Large Language Models

Q1: The ablation studies suggest most performance improvements stem from BDL.

A1: We need to clarify that DiffUp aims to reduce data and achieve efficiency. The BDL aims to improve performance. It is unfair to only compare accuracy improvement. To verify the effectiveness of DiffUp, we randomly select 55% data and distill with BDL. This setting requires the same number of iterations as DiffUp, and the results show the effectiveness of DiffUP.

Q2: Including more ablation studies would help clarify the contribution of each component.

A2: The two components of our DA-KD do not both contribute to performance, which has been discussed in Q1. To verify the contribution of each component, we conduct ablation studies on Llama2-1.3B.

We can see that DiffUp reduces the number of training iterations from 3570 to 1963. BDL improves the performance from 22.57 to 25.91. Combining DiffUP and BDL, we can achieve both efficiency and higher performance.

Q3: Including a discussion of how BDL compares to existing KL variants would strengthen this paper.

A3: Existing KL variants include two main types, mixing KL and RKL or mixing teacher and student distributions.

The former is formalized as a linear combination of KL and RKL [1]. It may suffer from gradient vanishing or explosion when facing extreme distributions. However, our BDL blends teacher and student distributions, avoiding the issues and achieving more stable updating (see Sec. 3.3 of the paper).

The latter category includes SKL [2]. However, it only mixes teacher into student distribution, which may cause gradient vanishing (see Eq. (6) in [2]). However, our BDL also mixes the student into teacher distribution, which can stabilize the gradient for better optimization (see Sec. 3.3 of the paper).

[1] Rethinking kullback-leibler divergence in knowledge distillation for large language models. (2024)

[2] Distillm: Towards streamlined distillation for large language models. (2024)

Q4: The overall novelty of this work is limited.

A4: Our DA-KD is designed to perform efficient LLM KD. We are happy that other reviewers recognize our contributions as "addresses an important bottleneck in LLM distillation", "simple yet effective" and "well-motivated and appropriate".

Q5: Have you evaluated BDL independently of DiffUp?

A5: We evaluate BDL independently and compare it with other losses. The results indicate that BDL achieves better performance.

Q6: Have you explored alternative sampling strategies?

A6: We have compared with random sampling in Tab 4 of our paper (denoted as w/o DDS). We also compare with the data filtered by the normalized DDS scores as probabilities. Results show stratified sampling performs better.

Q7: Have you examined the inclusion of samples generated by the teacher and the student?

A7: We incorporate teacher-generated or student-generated samples into the distillation. Here are the results. We can see a small performance degradation when using teacher-generated or student-generated samples.

We sincerely thank you for recognizing our contributions, by stating "well-motivated and appropriate", "experimental design is comprehensive" and "demonstrating reduced computational cost while maintaining/improving performance". Your thoughtful suggestions and questions help a lot to improve our work. We have carefully addressed your questions below:

Q1: The proposed BDL can be converted to KL (when $\lambda \rightarrow 0$) or RKL (when $\lambda \rightarrow 1$). How would it compare to combining KL and RKL linearly, i.e., $(1-\lambda)KL+\lambda RKL$?

A1: Thank you for your suggestion! We follow your suggestion and perform the experiments. Specifically, we use LLaMA2-7B as the teacher and use LLaMA2-1.3B as the student, and $\lambda$ is set to 0.9. The results are shown below. We observe our BDL performs better than the linear combination of KL and RKL. We hypothesize this is because BDL blends teacher and student distributions, avoiding the issues of gradient explosion or vanishing and achieving a more stable updating process, which is also explained in Sec. 3.3.

Q2： Is the cross-entropy (CE) loss incorporated when distillation is performed? If so, how is the ratio set?

A2：Thank you for your question. The cross-entropy loss is not used in distillation, because we find directly using KD loss performs the best performance. The experiments of loss function combined with cross-entropy loss are shown below, which is about 3~4% lower than ours. The similar observation was also found in Distillm[1].

[1] "Distillm: Towards streamlined distillation for large language models." arXiv preprint arXiv:2402.03898 (2024).

Q3: Discuss potential limitations of DA-KD.

A3: Thanks for your suggestion! One of the limitations of our DA-KD is that its performance is depended on the teacher model quality. If the teacher model provides incorrect prediction, the DDS mechanism may misidentify difficult samples, potentially affecting student performance. We will discuss these limitations in our final version and try to solve it in our future work.

We sincerely thank you for recognizing the strengths of our work by stating “simple yet effective”, "well-written", and "comprehensive experiments". Your thoughtful suggestions help a lot to improve our work. We have carefully addressed your concerns below:

Q1: DDS Computation & Dataset Partitioning (Weakness 1)

A1: The DDS computation and dataset partitioning occur before each epoch, dynamically updating difficulty scores based on the student model’s state. While it introduces extra computational cost, we can achieve a 55% reduction in training iterations. As a result, we can still achieve an overall reduction in computational cost compared to existing KD methods, which is shown in Table 3 in our paper.

Q2: Prior Work on Data Selection in KD (Weakness 2)

A2: Thanks for the suggestion and we apologize for the oversight in our discussion of the related works. Prior data selection methods are either solely based on student model [4] or solely based on the teacher model [5], ignoring their discrepancy. In contrast, our DDS metric integrates both teacher and student distributions. It selects samples that teacher can provide informative guidance to student. We will supplement a detailed discussion in our related work.

Q3：Generalisability of DDS to other domains (Weakness 3)

A3: To evaluate DDS in a different domain, we conduct experiments on CIFAR-100 using ResNet-50 as the teacher and ResNet-18 as the student. We compare three settings: (1) using the full dataset, (2) randomly selecting 50% of the data, and (3) selecting the top 50% of the data based on DDS. The results below demonstrate the effectiveness of DDS.

Q4: Comparison with SE-KD [3] (Questions 1)

[3] Song, Yuncheng, et al. "Self-Evolution Knowledge Distillation for LLM-based Machine Translation." 2024.

A4: SE-KD is conceptually related to our work by both proposing a difficulty metric of the samples and mixing the distributions of teacher and student. However, our method is intrinsically different from this work:

• Different objectives

  SE-KD aims to improve the LLM distillation performance, by dividing samples into two categories and learning with different metrics.

  However, DA-KD designs an efficient KD approach, which gradually reduces the data volume, and thus reduces the computational overhead.

• Different criterion for dividing samples

  SE-KD utilizes KL divergence between the student distribution and a target distribution to quantify the learning difficulty of samples. It does not consider whether a sample is well-learned by teacher.

  However, DA-KD uses the ratio between the student and teacher losses to measure the difficulty of samples. It only selects the samples that are well-learned by the teacher but poorly-learned by the student.

• Different distributions mixture in loss

  SE-KD mixes the student distribution and target distribution unidirectionally, which is similar to the SKL in Distillm[a].

  Differently, DA-KD achieves a bidirectional mixing of teacher and student distributions, leading to a more stable optimization process, which is explained in Sec. 3.3.

We perform an experimental comparison with SE-KD, using LLaMA2-7B (teacher) and LLaMA2-1.3B (student). The results are listed below. Our DA-KD outperforms SE-KD with less training time (about 16% reduction) while obtains better performance.

[a] "Distillm: Towards streamlined distillation for large language models." arXiv:2402.03898 (2024).

Q5: Theoretical Explanation of q >> p (Questions 2)

A5: In KD, the teacher model is well-learned and often produces a probability distribution that is smoother and more refined. Conversely, the student, being smaller or less trained, tends to be more overconfident in its incorrect predictions [b], or has less diversity in its output contents [c], resulting in a sharper distribution. This is characterized by high confidence concentrated distribution on fewer options. Therefore, for difficult samples, the probability of the wrong class predicted by student (represented by q) is significantly higher than that predicted by teacher (represented by p).

[b] "On student-teacher deviations in distillation: does it pay to disobey?." NeuIPS 2023.

[c] "Teaching large language models to reason with reinforcement learning." arXiv:2403.04642 (2024).

Q6: Several typos

A6: Thank you for pointing it out. We will correct all mentioned typos.

We sincerely thank you for recognizing the strengths of our work, including “addresses an important bottleneck in LLM distillation” and “experimental evidence supports the claims". Your suggestions are very helpful for improving our paper. We have carefully addressed your concerns below:

Q1: It would benefit from a discussion of alternative methods that focus specifically on dynamic dataset selection in KD. (Revisit the power of vanilla knowledge distillation: from small scale to large scale. Advances in Neural Information Processing Systems 36 (2023): 10170-10183.)

A1: Thank you for the suggestion. The suggested reference [1] aims to prove that VanillaKD can achieve comparable performance with other KD methods like DKD[2] and DIST[3] when using data augmentation and stronger training strategies instead of data selection. In contrast, we introduce a dynamic dataset selection mechanism to make the distillation process more efficient. We will add the discussion with the work [1] in our final version.

[1] "Revisit the power of vanilla knowledge distillation: from small scale to large scale." Advances in Neural Information Processing Systems 36 (2023): 10170-10183.

[2] "Decoupled knowledge distillation." Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition. 2022.

[3] "Knowledge distillation from a stronger teacher." Advances in Neural Information Processing Systems 35 (2022): 33716-33727.

Q2: During the Difficulty-aware Data Updating process, in order to obtain the difficulty scores for each sample, it is needed to calculate the cross-entropy loss of the teacher and student models separately for each sample, which may bring some computational budget.

A2: Although it will introduce extra computational cost when computing DDS, we can achieve a 55% reduction in training iterations by using DDS to filter the distillation data (Table 3 in the main paper). As a result, we can still achieve overall reduction in computational cost compared to existing KD methods.

To better analyze the efficiency of our DA-KD framework, we list the training time of different distillation methods when distilling Llama2-7B into a 2.7B. Our DA-KD requires only 26.1% of the training time compared to GKD (106.35 vs. 408.24 minutes) and 49.9% compared to Distillm (106.35 vs. 213.34 minutes), demonstrating the efficiency of the proposed method. We will discuss the data updating cost in our final version.

Q3: It would be helpful to include a discussion on potential limitations.

A3: Thanks for the suggestion! One of the limitations of our DA-KD is that its performance depends on the teacher model quality. If the teacher model provides incorrect answers, the DDS mechanism may misidentify difficult samples, potentially affecting student performance. We will discuss these limitations in our final version and try to solve it in our future work.

Q4: There are some typos or grammar errors: ... ...

A4: Thank you for pointing it out. We will correct all mentioned typos and proofread carefully.