MiniLLM: Knowledge Distillation of Large Language Models

The reviewer checked responses from the authors. The reviewer wants to thank them for their clarifications and updates.

We have replaced the GPT-4 score in Figure 1 and 5 with Rouge-L. We use GPT-4 for evaluation because it is studied by previous literature[1,2], where GPT-4 shows good alignment with human evaluation. This metric is also widely used in previous works to evaluate instruction-following[3,4,5].

It may be ok if the authors use the GPT-4 score (feedback) as a supplemental evaluation metric (i.e., moving them to appendix), but it is still not either convincing enough or clear how scientifically meaningful the evaluations are, for the following reasons.

• Being used in a few previous studies does not explain why the metric is scientifically meaningful for this study. [3] is a blog post, [4] failed to justify why GPT-4 feedback is meaningful, [5] analyzes "Inter-Annotator Agreement", but the protocol is unclear. Moreover, their assessment is based on labeling which response was better, or whether neither response was significantly better than the other, which is different from the scheme used in this study (score 1 - 10)
• GPT-4 is a black-box service, and its behaviors may change over time. Also, the service is updated in an untrackable manner, thus it is very challenging to reproduce experimental results based on such services (including [1] and [2]).

$w_t$ does not reduce the estimator's variance in Eq. 5. It is the important sampling weight to ensure that the gradient estimation by sampling from the teacher-mixed distribution equals the original estimator.

The reviewer needs more clarifications for this point. It looks like the revised manuscript still claims setting $w_t \approx$ ... to reduce the variance of the estimator in Eq. 5 as follows

Therefore, we approximately set $w_t \approx$ ... to reduce the variance of the estimator in Eq. 5

from the manuscript.

4. Two examples from the Unnatural-Instructions test set on which the 760M MiniLLM outperforms the 1.5B teacher model are as follows:

   Example 1

   ## Instruction:
   You are given a sequence of steps which describe how to perform some task. Each step might reference information from a previous step. Your job is to re-order the steps so that they can be performed in order and make sense.
   ## Input:
   [' Mix together flour, sugar, eggs and butter', ' Preheat oven to 375 degrees', ' Roll out dough on floured surface'].
   ## Ground Truth:
   [' Preheat oven to 375 degrees', ' Mix together flour, sugar, eggs and butter', ' Roll out dough on floured surface']
   ## Teacher Model:
   First, preheat oven to 374 degrees. Then, mix together flour, sugar, eggs, butter, and strawberry jam. Once the oven reaches the desired temperature of 374 degrees, you can begin preparing the delicious strawberry jam bars. Have a try!
   ## MiniLLM:
   1. Preheat oven to 375 degrees
   2. Mix together flour, sugar, eggs and butter
   3. Roll out dough on floured surface
   

   Example 2:

   ## Instruction:
   You will be given a paragraph of text, followed by a list of questions. For each question, you must determine the answer based on information in the paragraph.
   ## Input:
   Passage: The process of photosynthesis is used by plants to convert light energy into chemical energy that can be later used to fuel the plants' activities. This process involves two main steps: light absorption and carbon fixation. During light absorption,plant pigment molecules absorb photons and transfer their energy to chlorophyll molecules in the thylakoid membrane. In carbon fixation,the plant uses this excited state energy to convert CO2 gas into glucose sugar.
   Question 1: What are the two main steps of photosynthesis?
   Question 2: Where does light absorption take place?
   Question 3: What happens during carbon fixation?
   ## Ground Truth:
   1.The two main steps of photosynthesis are light absorption and carbon fixation.
   2. Light absorption takes place in the thylakoid membrane.
   3. During carbon fixation, the plant uses excited state energy to convert CO2 gas into glucose sugar.
   ## Teacher Model:
   Light absorption and carbon fixation. Light absorption takes place when chlorophyll molecules in plant cells capture photons from sunlight during the process of photosynthesis. This absorption of light energy is a crucial step in the conversion of solar energy into chemical energy…
   ## MiniLLM:
   Answer 1: Light absorption and carbon fixation. 
   Answer 2: Light absorption takes place in the chlorophyll molecule of the plant.
   Answer 3: In carbon fixation, the plant uses this excited state of energy to convert CO2 into glucose.
   

   From the examples, we observe that MiniLLM follows the instructions better when generating long responses. The teacher model follows the instructions at first but then generates contents unrelated to the instructions. We attribute this phenomenon to exposure bias in the teacher model[5]. During training, the teacher model predicts the next word conditioning on the prefix from real data distribution, but during inference, the prefix is from the teacher model’s distribution, introducing a discrepancy. And this discrepancy accumulates during generation, causing the responses to finally deviate from the correct answer (exposure bias).

   In contrast, MiniLLM sees self-generated responses during training (see Algorithm 1 in Appendix B), migrating the discrepancy. This effect is also observed in [6]. To summarize, MiniLLM with less exposure bias is possible to outperform the teacher model with more exposure bias.

5. We report the relative improvement when scaling up the teacher model compared to using the 340M teacher model. We find that MiniLLM benefits slightly better than SeqKD from scaling up the teacher model.

   Teacher Model Size	340M	760M	1.5B
   SeqKD	0.0%	0.3%	1.0%
   MiniLLM	0.0%	1.4%	3.4%

6. The value is Reverse KLD. We have fixed this typo in the revision.

7. The only specific hyper-parameter of MiniLLM is the $\alpha$ in Eq. 4, and we provide an ablation study in Appendix D.4. We find that $\alpha=0.2$ is generally suitable for models across sizes and larger models are more robust to the choice of $\alpha$.

[1] Defining and Characterizing Reward Hacking. 2022. In NeurIPS.

[2] A Survey on Model Compression for Large Language Models.

[3] Vicuna: An Open-Source Chatbot Impressing GPT-4 with 90%* ChatGPT Quality.

[4] Instruction Tuning with GPT-4.

[5] Why Exposure Bias Matters: An Imitation Learning Perspective of Error Accumulation in Language Generation. In ACL findings.

[6] Text generation by learning from demonstrations. 2021. In ICLR.

Regarding W2:

Point 1 & 4 (Novelty)

We study the white-box KD of LLMs on text generation tasks, but previous works on white-box KD focus on the image and text classification models, which differ largely from current LLMs. KD of LLMs has its own difficulty, and directly transferring previous white-box KD methods to LLMs does not work well (see the KD lines in our experiments). This is because LLMs generate answers auto-regressively, yielding a much larger label space than image and text classification tasks. Given the low capacity of the student model, there should be a trade-off between the mode-seeking and mode-average behavior. We point out this issue and provide a solution.

For the same reason, even if the properties of reverse v.s. forward KLD have been explored, applying reverse KLD to the KD of LLMs is still non-trivial, as stated in Section 2.2, and thus the three strategies are needed. Some of the strategies are inspired by prior work, but they are themselves novel:

• Single-step Decomposition: Prior work only shows the importance of the first-step reward, but we propose to decompose it from the total gradient.

• Teacher-mixed Sampling: The idea of mixing the teacher distribution is novel and the cited works are about analyzing the reward-hacking phenomenon, and the derivation of importance sampling.

• Length Regularization: This strategy is inspired by our pilot experiments and unrelated to prior works.

In summary, the contributions of our work are

• Discussion on the difficulty of knowledge distillation of LLMs, which is practical but far from solved.
• A novel KD objective for LLMs and a stable training approach.
• Extensive experiments on diverse datasets and model scaling property, which is essential for the practical use of the method on real-world LLMs.

Point 2 & 3 (Definition of white-box KD)

We follow the definition in [5], where black-box KD refers to the scenario where only the API of the teacher model is available (only the generated sentences, not including the output distribution)[6,7], and other cases are white-box KD. MiniLLM belongs to the white-box KD because we need the model parameters to compute output distribution. In other words, we study KD when the teachers’ output distribution is available, which is related to using reverse KLD for optimization. We have clarified the description of “white-box KD” in the introduction of the revised paper.

Regarding Question #2 (The student model outperforming the teacher model):

We suspect the reason is related to exposure bias, i.e., the discrepancy between the MLE training and text generation phases. Since the teacher model is fine-tuned with MLE on the instruction following dataset, it suffers from exposure bias. However, MiniLLM is trained with policy optimization, which mitigates exposure bias in language generation, verified by both prior work[8] and our experiments in Figure 6. Therefore, a small model with less exposure bias (MiniLLM) is entirely possible to outperform a large model with more exposure bias. The metrics we use are the common metrics to evaluate instruction-following (Rouge-L, GPT-4 feedback, human evaluation)[9,10]. We leave the study of how the choice of metrics related to this phenomenon to future work.

Regarding “doesn’t this graph (Figure 8) imply that w/o teacher-mixed sampling is better, if trained long enough?”

MiniLLM w/o Teacher-Mixed Sampling shows the reward hacking phenomenon[11], where the model obtains high reward (low reverse KLD) but low scores on the final evaluation metrics (Rouge-L and GPT4 feedback in Table 4).

[1] The Curious Case of Neural Text Degeneration. 2020. In ICLR.

[2] Tailoring Language Generation Models under Total Variation Distance. 2021. In ICLR.

[3] HONEST: Measuring Hurtful Sentence Completion in Language Models. 2022. In ACL.

[4] Training language models to follow instructions with human feedback. 2022. In NeurIPS.

[5] A Survey on Model Compression for Large Language Models.

[6] Vicuna: An Open-Source Chatbot Impressing GPT-4 with 90%* ChatGPT Quality.

[7] Instruction Tuning with GPT-4.

[8] Text generation by learning from demonstrations. 2021. In ICLR.

[9] Super-NaturalInstructions: Generalization via Declarative Instructions on 1600+ Tasks. 2022. In EMNLP.

[10] LIMA: Less Is More for Alignment. 2023. In NeurIPS.

[11] Defining and Characterizing Reward Hacking. 2022. In NeurIPS.

We thank the reviewer for their comments. Further responses and clarifications are as follows:

Regarding Novelty

We have clarified that "we study the property of reverse KLD for KD of LLMs" in section 2.1 and added relevant citations.

Regarding Diversity

• “Dist-4”[9] is a fraction: $N/C$, where $N$ is the number of the distinct 4-grams in the generated responses and $C$ is the total number of 4-grams. It is a widely used metric to measure the generation diversity of a language model.

• Loss: The loss we reported in the original response and Table 3 is the language modeling loss: $\text{Loss} = -\sum_{y\sim p_{\text{train}}} \log q_\theta(y|x)$. Therefore, it measures "how probable the instances in the training distribution are from the model distribution", rather than "how probable the generated text is from the training distribution". We have clarified this in the latest revision of our paper.

  In other words, loss measures the mode coverage of the real data distribution because it is essentially the forward KLD between the real data distribution and the model output distribution. This relates to diversity as in the ability to generate different generations given one context with different random seeds. Although this diversity is sacrificed to some extent (MiniLLM 3.95 v.s. SeqKD 3.91 in Table 3), we argue this is acceptable in practice because (1) the absolute value of the loss increase is small (2) many applications of test generation only require one but precise answer (like summarization, question answering, and coding).

The main point of the “Forward KLD v.s. Reverse KLD” part in our response is that using reverse KLD for MiniLLM
(1) will not lose the long-tail knowledge of LLMs (because we consider $p(y|x)$, not $p(x,y)$ and add $\mathcal{L}_{\text{PT}}$);

(2) will not introduce fairness and ethical issues (verified by the HONEST metrics);

(3) will probably lose the diversity of the generated sentences (because low-probability regions of $p(y|x)$ are not covered), but the diversity loss is not large (verified by the results in Table 3);

(4) is sufficient for KD of LLMs and simply combining forward KLD does not yield better performance.

In summary, we do not argue that coverage is in fact still intact (the Dist-4 indeed decreases and the loss indeed increases). We argue that using reverse KLD allows the MiniLLM to ignore modes that are small or even noisy low-probability regions of $p(y|x)$ (not $p(x,y)$) but still covering major (not all) parts of the teacher distribution, leading to a slight and acceptable drop in diversity and improvement in generation quality. This is sufficient for KD of LLMs in text generation.

White-box

In the context of KD for LLMs, it is a common practice to include the scenario where the output distributions are available in white-box KD[1,2,3], mainly to distinguish from the methods that only use the APIs (output distributions are generally unavailable in model APIs). If the reviewer finds it necessary, we can clarify the difference between the “division of white and black box” in KD for LLMs and adversarial attacks in the footnote to prevent misunderstanding from readers familiar with the adversarial attacks.

[1] A Survey on Model Compression for Large Language Models.

[2] Lion: Adversarial Distillation of Closed-Source Large Language Model.

[3] LaMini-LM: A Diverse Herd of Distilled Models from Large-Scale Instructions.