DDK: Distilling Domain Knowledge for Efficient Large Language Models

Summary

This paper proposes DDK, a knowledge distillation (KD) framework that distills large language models (LMs) into small LMs. Unlike previous KD methods, DDK dynamically adjusts the domain weights during distillation. Experiments show that DDK outperforms other KD baselines across various tasks.

Strengths

1. The paper is well written and the method is easy to follow. 2. The experiments show that DDK outperforms other KD baselines on various tasks.

Weaknesses

The extra computation introduced by KKD should be considered. It seems KKD requires the inference of a large LM during the training of the small LM. When the teacher model is much larger than the student model (QWen-1.5 14B v.s. QWen-1.5 1.8B), the inference cost of the teacher model would be even larger than training the student model. Therefore, it is more reasonable to compare the performance of the distilled model and the baselines given the same FLOPs.

Questions

1. What are the training data for the baselines (CPT, TED, KD, and MiniLLM)? Is the data for DDK the same as that for the baseline methods? 2. In lines 178-179, is the learning rate 3e-5 ($3\times 10^{-5}$) rather than $3e^{-5}$?

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

N/A

Summary

The paper introduces a new framework called Dynamic Domain Knowledge Distillation (DDK) to enhance the efficiency of knowledge distillation for large language models (LLMs). Unlike traditional methods that overlook domain performance differences between student and teacher models, DDK dynamically adjusts the distillation dataset composition based on these differences, ensuring a more stable and effective knowledge transfer. This approach addresses the issue of excessive focus on domains with minimal performance gaps and enhances overall model performance. Extensive evaluations demonstrate that DDK significantly outperforms existing knowledge distillation methods and continuously pretrained baselines.

Strengths

- The proposed dynamic dataloader for KD is technically sound. - Numerical experiments well validate the efficacy of the method.

Weaknesses

- Dynamic dataloader requires knowing the training data distribution and category beforehand. - Missing references. Similar ideas have been explored in pruning LLMs, such as ShearedLLaMA, LoRAShear to recover the knowledge . The paper needs to discuss with them in the related work section due to the closed relation between pruning and KD. Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning LoRAShear: Efficient Large Language Model Structured Pruning and Knowledge Recovery

Questions

- How the methods perform on under other KD losses, such as reversed KLD, JSD, skew-KLD. On-policy distillation of language models: Learning from self-generated mistakes. DistiLLM: Towards Streamlined Distillation for Large Language Models.

Rating

6

Confidence

5

Soundness

3

Presentation

3

Contribution

2

Limitations

Yes.

Summary

The work introduces a novel framework for knowledge distillation (KD) for LLMs. The key innovation of DDK is its dynamic adjustment of the distillation dataset composition based on domain performance differences between the teacher and student models. The paper presents extensive evaluations demonstrating that DDK significantly improves performance in various KD settings, outperforming both continuous training baselines and existing KD methods.

Strengths

1. The authors provide extensive empirical evidence demonstrating the effectiveness of DDK in improving the performance of student models across various benchmarks. 2. As the computational and storage demands of LLMs are significant barriers to their widespread deployment, KD is a promising solution. The proposed KD method is simple and easy to follow.

Weaknesses

1. Discuss the difference between DDK and the Dynamic Batch Loading proposed by Sheared LLaMA[1], which is also proposed to adjust domain proportions for dynamically training smaller models. They also identify discrepancies in loss between smaller and larger models across various domains, and accordingly, they sample more data from domains where the discrepancy is more pronounced. While they concentrate on structural pruning, it is akin to the DDK. Consequently, I perceive the novelty of DDK as being somewhat limited. 2. The results of Qwen 1.5 in Table 1 are not significantly convincing. The MMLU/HumanEval of Qwen 1.5 1.8B in the Qwen official blog are 46.8/20.1 while the authors' report is 44.5/11.9. In addition, compared to the official results, we can see that the DDK fails to improve the model of the students on MMLU. The authors need to check this and provide **more robust results of baselines**. [1] SHEARED LLAMA: ACCELERATING LANGUAGE MODEL PRE-TRAINING VIA STRUCTURED PRUNING. Xia et al., 2023

Questions

1. In the setting of paper, the domains are predefined. How to extend the DDK framework for new domain during the distillation training process? Could give more experiments on continual domain learning settings?

Rating

5

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

Yes.

Summary

This work proposed a KD strategy for LLMs. Specifically, with assess to the domain-specific performance of both the teacher and student LLMs, DDK uses domain knowledge guided sampling to dynamically update the data mixture. In addition the paper also conducts a statistical analysis of the domain distribution of the datasets involved. The training process is relatively straightforward and easy to generalize. The experimental results also show that DDK's training method improves the performance average of different data sets.

Strengths

1. A complete training algorithm is designed, and the process is explained clearly. The process of the DDK algorithm is easy to extend to the training process of other models. 2. The authors conducted a comprehensive knowledge distillation experiment on two large model families and a comprehensive ablation study.

Weaknesses

1. Although the method proposed in this paper is easy to understand and effective, I doubt that the method in this paper is limited to LLMs. In other words, this paper does not mention (or needs to explain) how previous researchers (before LLMs) performed domain-enhanced distillation for domain-biased datasets, and why these previous methods cannot be applied to the distillation of LLMs to achieve similar results. The advantages and novelty of this paper's domain sampling method over previous work that may be transferable to LLMs need further explanation. 2. In the experimental part, there is a lack of key comparison between DDK and other methods that focus on similar domain sampling. The baseline actually involves the work that focuses on domain in KD (cited as [60], etc.), but the subsequent analysis only compares the total average score of DDK and these works, which seems to lack comparison and analysis of similar works. As far as I know, other baselines are more general KDs, and do not focus on domain information. It is certainly worth noting that DDK performs better than baselines such as MiniLLM, but I think what can better illustrate the effectiveness and novelty of this paper is the comparison with similar domain data sampling, including experimental analysis. 3. In the experimental section, you can add experiments on the dataset and the scale property of the teacher model. This is a possible suggestion.

Questions

The questions I expect to ask would be similar to the above section.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The method proposed in this paper is very complete and solves the problem well, but there are perhaps two points to note: 1. The distillation method designed in this paper does not seem to be necessarily limited to LLMs. One of the main difficulties of KD on LLM may be that the distribution difference between the teacher and student models is too large, but the key points that this paper focuses on and solves seem to be orthogonal to this point. So it also leads to a similar question: why previous similar methods that focus on domain sampling cannot be migrated to LLMs, and what are the advantages and novelties of this paper's design. 2. Following point 1, in the experiment, what are the specific advantages of this method over the predecessors in the domain problem (not just the overall average, of which result shows DDK outperformed others). This may be what I am very curious about after reading the paper.

Official Comment by Reviewer 6kDi

I thank the authors for their response. After reading the response, I think my current score is appropriate.

Summarization on the Responses

Thanks for handling/reviewing our submitted manuscript: "DDK: Distilling Domain Knowledge for Efficient Large Language Models". We would like to thank the reviewers for their insightful and constructive comments and suggestions. By addressing each of the issues raised by the reviewers, we believe that the quality and clarity of our DDK can be improved a lot. The major responses are summarized as follows: (1). We have carefully discussed the advantages and novelties of DDK (See Reviewer foae.Q3, Reviewer EYh3.Q1) (2). We have additionally compared the domain-enhanced KD methods and discussed the advantages over previous methods in detail (See Reviewer foae.Q1&Q4). (3). We have provided experiments on another training dataset and discussed the scale property of teacher (See Reviewer foae.Q2). (4). We have provided more clarification and discussion on the evaluation prompts (MMLU and Humeneval) and continual domain learning (See Reviewer EYh3.Q2&Q3&Q4) (5). We have discussed the limitations of requiring data distribution and provided experiments based on the clustering method to address this limitation (See Reviewer oZHE.Q1). (6). We have discussed more related works (e.g., Sheared LLaMA, LoRAShear, Skew-KLD) and conducted a detailed analysis of these works (See Reviewer EYh3.Q1, Reviewer oZHE.Q2&Q3). (7). We have provided a more detailed analysis and clarification on the training costs, training data, and learning rate (See Reviewer 6kDi.Q1&Q2&Q3). Again, we would like to sincerely thank you very much for these constructive comments and evaluation for our manuscript.

Looking forward to feedback on the Responses

Dear Reviewers: Hello! Since the discussion period has started, we would like to kindly ask you to take a look at our responses and reevaluate our work based on our clarifications. Please let us know whether our response addresses your concerns or whether there is any further detail we can provide to help address them. We appreciate your time and consideration!

Looking forward to Feedback as Discussion Deadline Approaches

Hi, we sincerely thank you very much for these constructive comments and evaluation of our manuscript. As the discussion phase will be closed on Aug 13 11:59 pm AoE, we would like to kindly ask you to take a look at our responses and reevaluate our work based on our clarifications. Please let us know whether our response addresses your concerns or whether there is any further detail we can provide to help address these concerns. Thank you again for dedicating your time to reviewing our paper.

Looking forward to Feedback as Discussion Deadline Approaches

Hi, we sincerely thank you very much for these constructive comments and evaluation of our manuscript. As the discussion phase will be closed on Aug 13 11:59 pm AoE, we would like to kindly ask you to take a look at our responses and reevaluate our work based on our clarifications. Please let us know whether our response addresses your concerns or whether there is any further detail we can provide to help address these concerns. Thank you again for dedicating your time to reviewing our paper.

Combine MiniLLM with the domain sampling method of DDK.

The comparison with MiniLLM in a similar domain sampling strategy is also discussed in **Q3** from **Reviewer oZHE**, where we can directly replace the KL divergence loss in Eq. 3 with the reversed KLD loss in MiniLLM. In **A3** for **Q3** from **Reviewer oZHE**, we call this alternative method as **DDK (Reversed KLD in MiniLLM)**, and we observe that replacing the distillation loss does not bring more gains when compared with the DDK with default KL divergence loss. For this phenomenon, just discussed in **A3** , we assume that our DDK focuses on the pre-training setting with a relatively large training data set, where the data quality and data mixture are very important for pretraining (See **G.Q1**). Meanwhile, the MiniLLM focuses on the SFT phase to address the mode-averaging problem in distillation. Therefore, we assume that the main contribution on the improvements comes from changing the domain mixture in our DDK. We will include this discussion in our new version.