DDK: Distilling Domain Knowledge for Efficient Large Language Models

Thanks for your careful reading and constructive suggestions. We will address your concerns shown below in detail.

Q1: Extra computation introduced by DDK should be considered. Compare the performance of the distilled model and the baselines given the same FLOPs.

A1: Thanks for your insightful suggestions. In Table 13 of Appendix B.2, we also report the training costs (i.e., TFLOPs) of different baseline methods, and we observe that the additional computation costs are acceptable when compared with the baseline KD method. Besides, we observe the ratio between the TFLOPs of DDK and TFLOPs of CPT is (5.401e8/1.456e8)=3.709. Then, we also perform the experiments by continuing pre-training the student model with about 56B tokens, where the TFLOPs is about 5.4e8. The results are shown in the following table. In the table, we observe that when we continue pretraining the student model for more tokens without using the teacher guidance, little gain is obtained. We assume that the Qwen1.5 model has been trained on >3T tokens, and the convergence status is stable when using the standard pretraining method. In contrast, when using the teacher guidance of the DDK, we observe consistent performance improvements in Table 1, Table 2, Table 7, and Table 8, which further demonstrates the effectiveness of our DDK.

We will revise the above discussion in our new version.

Q2: Is the data for DDK the same as that for the baseline methods?

A2: The training data is the same and we adopt the same seed for shuffling the dataset, which can ensure the training data same for different methods. We will clarify this detail in our new version.

Q3: Learning rate issue.

A3: Thanks for your kind suggestions. We will revise this typo in the new version.

Thanks for your careful reading and constructive suggestions. We will address your concerns shown below in detail.

Q1: Require knowing the training data distribution and category beforehand.

A1: Our DDK requires the training data distribution and category to divide the domains and we acknowledge this limitation. However, it should be mentioned that existing pretraining datasets usually provide the source composition. For example, the Repajama[https://www.together.ai/blog/redpajama] includes 7 domains (i.e., CommonCrawl, C4, Github, Wikipedia, Books, ArXiv, and StackExchange), The Pile [https://arxiv.org/pdf/2101.00027] includes 22 domains (i.e., Pile-CC, PubMed Central, Books3, OpenWebText2 and etc.) and Dolma[https://arxiv.org/pdf/2402.00159] includes 6 domains (i.e., Common Crawl, GitHub, Reddit, Semantic Scholar and etc.)

Meanwhile, we also tried another option to dive the training data into 10 domains based on the clustering method to address this limitation, where we first extract the feature of each document using [https://huggingface.co/BAAI/bge-m3] and then use k-means to divide the training corpus into 10 domains implicitly. Then, we apply DDK on the divided implicit domains based on the clustering and we call this method DDK (clustering), where we do not need to know the training data distribution and category. In the following table, we observe that the DDK (clustering) achieves relatively comparable performance with original DDK. We assume that the explicitly divided domains based on the sources (e.g., The stack, OpenWebMath) are distinct. Besides, we will continue to investigate the clustering method, which is more scalable.

Q2: Cite and discuss shearllama and lora shear.

A2: Please See General Response (G.Q2) on the discussion between Sheared LLaMA and DDK.

The discussion between LoRAShear and DDK is as follows.

The LoraShear is proposed for structured pruning on LLMs, which first creates the dependency graphs over LoRA module and then proceeds progressive structured pruning on LoRA adaptors. To recover the lost knowledge during pruning, LoRAShear proposes a Dynamic Knowledge Recovery scheme on both the pre-training and instructed fine-tuning datasets to effectively narrow down the performance gap to the full models. In contrast to LoraShear, DDK aims to efficiently transfer the domain knowledge of the teacher network, and propose a factor smooth updating strategy to stabilize the domain knowledge guided sampling process.

We will cite these works in our new version.

Q3: How the methods perform under other KD losses (reversed KLD, JSD, skew-KLD)?

A3: To address the mode-averaging problem of the student model, where the student model learns an overly smooth distribution in an attempt to cover the teacher’s entire support set, the reversed KLD in MiniLLM and generalized GSD in [R1] are proposed. However, these on-policy approaches (i.e., the reversed KLD and the generalized JSD) show lower efficiency, the skew-KLD in [R2] introduces the adaptive off-policy approach to enhance the efficiency. In contrast to these works, our DDK focuses on efficiently transferring the domain knowledge of the teacher network based on standard forward KD loss, where the motivation and technical details are different from these works a lot. Besides, our DDK is also orthogonal to these KD losses, where we can directly replace the KL divergence in Eq. 3 with the above losses in MiniLLM, [R1] and [R2] easily. In the following table, we report the results of our DDL using different distillation losses, and we observe that these losses have not provided more performance gains for our DDK based on KLD. For this phenomenon, 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, in MiniLLM and [R1, R2], these losses mainly focus on the SFT phase to improve learning efficiency by addressing the mode-averaging problem. Therefore, we assume that the main contribution on the improvements comes from changing the domain mixture in our DDK. In the future, we will also continue to investigate the optimal distillation loss formulation for our DDK following [R1, R2].

[R1]: On-policy distillation of language models: Learning from self-generated mistakes, https://arxiv.org/pdf/2306.13649

[R2]: DistiLLM: Towards Streamlined Distillation for Large Language Models, https://arxiv.org/pdf/2402.03898

Thank you for your valuable comments.

Q1: Discuss with Sheared LLaMA. Novelty.

A1: Please See General Response (G.Q1 and G.Q2).

Q2: Qwen1.5 results on MMLU and Humaneval.

A2: For MMLU, the accuracy from Qwen Blog is based on a 5-shot setting, while the result in Table 1 is based on a zero-shot setting (See Line 182). Besides, we provide the 5-shot result of MMLU in Table 12 and observe the result is 45.59, which is close to the provided result (46.8) of the Qwen blog. Besides, we assume the gap between 46.8 and 45.59 comes from the evaluation prompt. Specifically, the evaluation prompt for DDK is as follows:

# Prompt-1
for opt in ['A', 'B', 'C', 'D']:
    prompt="Answer the question.\nQuestion: {question}\nChoices:\nA. {A}\nB. {B}\nC. {C}\nD. {D}\nAnswer: " + opt

The {question}, {A}, {B}, {C}, {D} are the corresponding inputs for each evaluation sample.

We also have tried the following prompt from OpenCompass [https://github.com/open-compass/opencompass], where the task description is also added to the evaluation prompt. Note that the mmlu_all_sets contains 57 tasks. Based on this prompt, we obtain 46.88 and 47.05 under the zero-shot and few-shot settings, respectively.

# Prompt-2
mmlu_all_sets = ['college biology', 'college chemistry', 'college computer science', ..., 'human aging', 'us foreign policy', 'conceptual physics']
for task in mmlu_all_sets:
    for opt in ['A', 'B', 'C', 'D']:
        prompt="The following are multiple choice questions (with answers) about" + task + '\n\n' + "{question}\nA. {A}\nB. {B}\nC. {C}\nD. {D}" + '\n' + 'Answer: ' + opt

We also provide the following results and DDK achieves consistent improvements in MMLU.

For Humeneval, we observe that the results of Qwen1.5-1.8B are very sensitive to the prompt, and we use the naive testing prompt (accuracy of 11.87) as follows:

# Prompt-1
prompt='{prompt}'

Similarly, we have tried following two prompts following OpenCompass:

# Prompt-2
prompt='Complete the following python code:\n{prompt}'
# Prompt-3
prompt='You are an intelligent programming assistant to produce Python algorithmic solutions.\nCan you complete the following Python function?\n```python\n{prompt}\n```'

And we observe the results of Prompt-2 and Prompt-3 are 23.17 and 26.83, respectively, which are higher than the reported result (20.1) on Qwen blog.

Meanwhile, we report the DDK results on these three evaluation prompt settings. In the following table, we observe that DDK still achieves better performance than the baseline. In new version, we will provide a detailed analysis of the evaluation prompt and provide more solid and robust results.

Q3: Extend DDK for new domains during distillation.

A3: In distillation, when a new domain is introduced, we first need to prepare the validation dataset for this domain following Appendix B.1 (Line 490-510). Then, following Eq. (1), based on the current model checkpoint, we compute the perplexity scores over the validation sets of all domains for student and teacher respectively, and obtain the newly initialized domain discrepancy factor, which is then used to change the domain mixture following the Algorithm 1 of the main paper.

Q4: Experiments on continual domain learning setting.

Q4: We provide the results of involving a new source domain data (i.e., Law) using the Pile-of-Law dataset [https://huggingface.co/datasets/pile-of-law/pile-of-law]. Specifically, we first train the original 10 domains for 5B/10B tokens. Then, we append the Law domain into the training dataset. After that, following Algorithm 1 of the main paper, we continue the distillation training on the student for additional 10B/5B tokens on these 11 domains, which aims to achieve a fair comparison with the DDK baseline results on 10 domains for about 15B tokens. To show the effect of introducing the law domain, we additionally report the results on the law subset of MMLU-Pro[https://huggingface.co/datasets/TIGER-Lab/MMLU-Pro]. (MMLU-Pro has 14 high-quality subsets with challenging questions.), and observe that by introducing the law domain, better results on the MMLU-Pro(Law) are obtained, and results on other datasets are also preserved well.

Thank you for your nice comments and suggestions.

Q1: Compare domain-enhanced KD methods.

A1: After investigating KD [R1] and applications on LLMs [R2], we have observed existing domain-enhanced KD methods can be divided into two categories. The first is cross-domain KD for domain adaptation to improve the generalization abilities of the unseen target domains based on seen source domains. In contrast, DDK aims to improve training efficiency and effectiveness given multiple seen domains, and these cross-domain KD methods cannot be applied to LLM pretraining. The second is based on multi-teacher strategy [R3, R4], which first trains different domain expert models and then distills the overall domain capacities into one single model. Thus, we implement two baselines following [R3, R4].

[R1]. Knowledge Distillation: A Survey

[R2]. A Survey on Knowledge Distillation of Large Language Models

[R3]. BAM! Born-Again Multi-Task Networks for Natural Language Understanding

[R4]. Knowledge Fusion of Large Language Models

Specifically, first, we train the domain expert models on 10 different domains as teacher models, where each domain expert is trained on about 5B tokens on the corresponding domain. Then, for the first baseline [R3], we select a corresponding teacher to produce logits for each sample based on the belonged domain. For the second baseline [R4], we ensemble logits from different domain teacher models. The table below lists the PPL on validation datasets for 10 domains.

We also report the downstream results and observe that baseline methods are inferior to DDK. We assume that these methods only consider producing better teacher guidance, and do not consider the effect of the domain mixture for LLM pretraining.

Q2: Experiments on dataset and scale property of teacher.

A2: For experiments on dataset, we have performed experiments using StarCoder on Stack V2 dataset in Table 11. Besides, we perform experiments on The Pile dataset [https://arxiv.org/pdf/2101.00027] with 22 domains, and we use Qwen-1.5 14B and Qwen-1.5 1.8B as teacher and student following Table 1. Note that as The Pile focuses on English, we do not evaluate CEval and C3.

We observe that DDK achieves better performance than CPT baseline on The Pile and results on The Pile are relatively lower than results in our main paper. We assume that the data quality used in our main paper (See Line 166-171) is better than The Pile, as we directly use the original Pile dataset in 2021 without additional cleaning strategies.

For scale property of teacher, we have used different sizes of teacher models (i.e., Qwen1.5-14B and Qwen1.5-7B) to distill Qwen1.5-1.8B in Table 1 and Table 8, and average results distilled by Qwen1.5-14B and Qwen1.5-7B are 55.41 and 53.64, respectively, which means that better results are achieved using a relatively larger teacher.

Q3: Advantages and novelties of DDK.

A3: Please See General Response (G.Q1).

Q4: Advantages over predecessors (not the overall average).

A4: First, in Fig. 1, DDK can reduce the performance gaps between teacher and student a lot in many domains. For example, for DDK, PPL on the Books domain is less than other methods a lot, and PPL on StackExchange domain is very close to teacher.

Second, in Table 1, we observe that a large relative performance gap usually exists in reasoning-related tasks for teacher and student. For example, relative performance gap ratios are 43.22% [(67.63-38.4)/67.63=43.22%] and 68.59%[(37.80-11.87)/37.80=68.59%] for GSM8K and Humeneval, respectively. In contrast, for knowledge-related tasks, the relative performance gap ratios are 24.17% [(78.68-59.66)/78.68=24.17%] and 30.87%[(64.34-44.48)/64.34=30.87%] for CEval and MMLU. After using DDK, the relative performance gaps are 20.84%, 28.23%, 18.98%, 28.49% on GSM8K, Humeneval, CEval and MMLU. In contrast, for ``CPT & DoReMi'' and MiniLLM, the relative performance gaps are [32.18%, 76.85%, 21.91%, 30.15%] and[27.69%, 55.34%, 21.63%, 29.95%], respectively. Thus, we assume that DDK can greatly improve the weaknesses on reasoning-related tasks.