MiniLLM: Knowledge Distillation of Large Language Models (2306.08543v4)

Abstract: Knowledge Distillation (KD) is a promising technique for reducing the high computational demand of LLMs. However, previous KD methods are primarily applied to white-box classification models or training small models to imitate black-box model APIs like ChatGPT. How to effectively distill the knowledge of white-box LLMs into small models is still under-explored, which becomes more important with the prosperity of open-source LLMs. In this work, we propose a KD approach that distills LLMs into smaller LLMs. We first replace the forward Kullback-Leibler divergence (KLD) objective in the standard KD approaches with reverse KLD, which is more suitable for KD on generative LLMs, to prevent the student model from overestimating the low-probability regions of the teacher distribution. Then, we derive an effective optimization approach to learn this objective. The student models are named MiniLLM. Extensive experiments in the instruction-following setting show that MiniLLM generates more precise responses with higher overall quality, lower exposure bias, better calibration, and higher long-text generation performance than the baselines. Our method is scalable for different model families with 120M to 13B parameters. Our code, data, and model checkpoints can be found in https://github.com/microsoft/LMOps/tree/main/miniLLM.

• The paper introduces a novel reverse KLD approach for knowledge distillation that mitigates overestimation of low-probability outputs in generative tasks.
• It employs optimization techniques such as single-step decomposition and teacher-mixed sampling to stabilize training and improve model convergence.
• The method scales from 120M to 13B parameters, achieving superior performance and calibration by reducing exposure bias in student models.

MiniLLM: Knowledge Distillation of LLMs

The paper "MiniLLM: Knowledge Distillation of LLMs" introduces a novel approach to distilling LLMs into smaller, more efficient ones through a method known as knowledge distillation (KD). The work specifically targets the need for more effective strategies in white-box KD, where the teacher model's distributional outputs are leveraged for distillation, aiming to overcome the limitations of existing techniques primarily used in classification tasks or with black-box model APIs.

Reverse Kullback-Leibler Divergence

The central innovation in this research is the replacement of the traditional forward Kullback-Leibler divergence (KLD) with reverse KLD as the objective for the KD process. This shift addresses the problem of overestimating low-probability regions of the teacher distribution, a common issue when using forward KLD in generative tasks. Reverse KLD encourages the student model to focus on key modes of the target distribution, thus optimizing model performance in core areas without unnecessary complexity.

Figure 1: Comparison between sequence-level KD (left) and MiniLLM (right). Sequence-level KD forces the student to memorize all samples generated by the teacher model, while MiniLLM improves its generated texts with the teacher model's feedback.

Optimization and Algorithmic Enhancements

To efficiently implement the reverse KLD objective, the authors propose several optimization strategies. The approach involves policy gradient methods for deriving the objective's gradient, facilitating precise learning of the intended distribution. Key strategies include:

• Single-Step Decomposition: Reduces training variance by focusing on individual generation quality at each step.
• Teacher-Mixed Sampling: Mitigates reward hacking by blending teacher and student distributions during sampling.
• Length Normalization: Corrects for biases towards short outputs by normalizing reward accumulations over sequence length.

These contributions stabilize training and improve convergence, ensuring the student model learns efficiently from the teacher's guidance.

Scalability and Performance

The experiments demonstrate that MiniLLM consistently delivers superior performance across models of varying sizes, from 120M to 13B parameters, within instruction-following tasks. The method's scalability is evidenced by its ability to maintain or exceed teacher-level performance. This is especially notable in terms of precision, measured by Rouge-L scores, where MiniLLM sometimes even surpasses its teacher models, thanks to reduced exposure biases.

Figure 2: The scaling law of teacher model based on the OPT family models. We compare MiniLLM and SeqKD with OPT-1.3M as the student and OPT 2.7B, 6.7B, and 13B as teachers.

Implications for Calibration and Exposure Bias

MiniLLM's focus on optimizing reverse KLD results in better-calibrated models with lower exposure bias. Calibration improvements are significant, particularly in avoiding the pitfalls of standard KD, which often distort probability distributions due to forward KLD. Moreover, the research suggests a promising direction towards effectively calibrating student models without sacrificing the generative diversity necessary for robust NLP applications.

Figure 3: The excess error caused by the training-decoding discrepancy (ExAccErr) accumulated with the generation length. Lower ExAccErr means the method introduces less exposure bias.

Conclusion

MiniLLM represents a significant advancement in the field of KD for LLMs, providing a scalable, efficient approach to reduce the computational requirements of deploying large models. By rethinking the divergence metric and optimizing the training procedure, the methodology not only enhances performance and scalability but also addresses fundamental challenges in model calibration and bias. This work sets the stage for broader applications of distillation techniques, potentially influencing future AI systems' design and deployment strategies.