GOLD : Generalized Knowledge Distillation via Out-of-Distribution-Guided Language Data Generation (2024)

Abstract: Knowledge distillation from LLMs is essential for the efficient deployment of language models. Prior works have proposed data generation using LLMs for preparing distilled models. We argue that generating data with LLMs is prone to sampling mainly from the center of original content distribution. This limitation hinders the distilled model from learning the true underlying data distribution and to forget the tails of the distributions (samples with lower probability). To this end, we propose GOLD, a task-agnostic data generation and knowledge distillation framework, which employs an iterative out-of-distribution-guided feedback mechanism for the LLM. As a result, the generated data improves the generalizability of distilled models. An energy-based OOD evaluation approach is also introduced to deal with noisy generated data. Our extensive experiments on 10 different classification and sequence-to-sequence tasks in NLP show that GOLD respectively outperforms prior arts and the LLM with an average improvement of 5% and 14%. We will also show that the proposed method is applicable to less explored and novel tasks. Code is available in the Appendix.

## introduction large language models (llms) have shown outstanding few-shot performance in solving different complex natural language tasks @xcite . the term few-shot refers to the ability of the llm to understand and perform tasks accurately given only a few examples. however, achieving such performance necessitates models with a large number of parameters to generalize and learn distinct tasks. the computational complexity of llms hinders their real-world applications for deployment. moreover, most of the llms are not publicly available and users should share their confidential data with llms through prompting which is a privacy concern. therefore, user-specific small models are critical to address efficiency and privacy concerns. to circumvent the above challenges, knowledge distillation (kd) from llms has been used to prepare small language models (slms). there are two main paradigms in kd: data-informed and datafree methods @xcite @xcite . data-informed methods are the conventional kd techniques that use llms to label data and train a small language model with the labels @xcite . data-free methods, on the other hand, study an extreme case when no human-generated dataset (unlabeled or labeled) is available during distillation. zerogen @xcite and progen @xcite are two prior works that proposed data-free knowledge distillation. zerogen @xcite proposed generating data for classification tasks using llms and trained a small model on the generated data. progen @xcite improved the quality of the generated data by finding important samples via an influence function in the course of data generation. these works are not able to find failure modes of slms and steer data generation toward the samples that can improve the generalizability of the model. @xcite showed that distilling knowledge from llms via data generation causes irreversible defects in the slm, where the tails of the original content distribution disappear. llms usually tend to generate samples with higher likelihood (known as in-distribution samples) repeatedly @xcite which results in poor generalizability of the distilled slm. preserving the llm's ability to model low-probability events or out-of-distribution (ood) samples is essential to the fairness of their predictions. such events are often relevant to marginalized groups and are vital to understand complex systems @xcite . usually, the performance of language models improves when more human-labeled data is available, since a larger dataset covers wider task domains. on the contrary, @xcite shows that increasing the amount of generated data with llms lowers the performance of the distilled slm. we argue that vanilla data generation with llms is prone to only generate high likelihood samples which negatively affects the slm performance with increasing the number of samples. to this end, we propose gold, a data-free kd framework that iteratively finds failure modes of the slm (i.e., ood samples) and provides feedback to the llm (figure 1 ) for the next iteration of data generation. figure 2 shows the overall framework of gold. the user provides the task definition with a few samples of the data corresponding to that task. the llm is then used to generate a training batch of data for the specified task, which is used to update the weights of the slm. given the generated train batch, we use prompting to ask llm to generate a separate batch of ood data that is significantly different from the train batch in terms of topic and style this ood batch is then used as our validation set to evaluate the performance of the slm and accordingly identify its failure modes as a feedback to the llm. we use the output logits of the slm to measure the energy values of the validation samples and select the top ones based on their corresponding free energy scores @xcite . the energy score does not require the labels of the generated sample and therefore is not prone to pick the data with noisy labels as the ood samples. the selected ood samples are then used as the feedback to the llm for the next iteration of data generation. the major contributions of this paper are: limitations of data generation with llms. two recent works @xcite challenge the idea that training a generative model on the output of a different generative model will lead to performance gains in the case of a smaller "imitation models" being trained to match the performance of a "teacher model" as in @xcite . however, we note their negative findings only apply in the most general case, where the synthetic train data is not "corrected" after being generated, either algorithmically or manually by humans before being hosted onto the web. further, the future-generation models considered in both works are not "task-specific" -they are trained with the goal of either meeting or exceeding the performance of the previous model, and the negative findings apply only in this case. our goals are less ambitious, and tailored towards achieving the performance of the teacher model (i.e., an llm) only on a specific task, rather than training a student model (i.e., an slm) to reproduce the entire suite of capabilities as the llm. in addition, @xcite indicates the poor generalizability of the distilled slm due to the generation of only high likelihood data by the llm. however, unlike the previous works, gold can effectively address this issue by iterative oodbased feedbacks to the llm to identify the failure modes of the slm and improve its generalizability. ## method fig. 2 shows an overview of the proposed gold, a data generation and kd framework for language models. we iteratively transfer knowledge from an llm, denoted by m l , to an slm, denoted by m s , via the data generated by m l . the kd is performed for any task that falls within the realm of the expertise of m l . the framework operates with two inputs: a task definition p task and a few samples of real data x real . in each iteration, gold generates a train set and an ood validation set denoted by (x train , x val ) that is respectively used to update the weights of m s and to evaluate the updated m s . a feedback function, denoted by e, is then used to find the top ood samples from x val that are within a pre-defined upper and lower bound threshold. the sample selection procedure is performed regardless of the generated labels, which makes the feedback function robust to noisy labels. the selected samples, denoted by x fb , are concatenated to the prompt for the next iteration of data generation. in the following, the details of the problem formulation, data generation procedure, the feedback function, and slm training are discussed. ## experiments in this section, the performance of gold is analyzed and compared with previous works over different classification and sequence-to-sequence tasks. first, we describe the experimental settings including the datasets (tasks), model architectures, prompts, and baselines. following that, the experimental results are quantitatively and qualitatively discussed and compared with the baselines. an extensive set of ablations over gold components, size of llm and slm, and the number of generated samples is also provided. ## conclusion in this paper, we proposed a task-agnostic framework for data generation and kd from llms to slms that is applicable to any nlp task, including new tasks. in order to improve the generalizability of the distilled slms, we introduced an iterative ood-empowered feedback mechanism. an energy-based ood evaluation approach was also proposed to handle noisy data generated by llms. with an extensive set of experiments, we showed that our method achieved state-of-the-art results on a variety of classification and sequenceto-sequence downstream tasks in nlp. future direction of this work includes applying the same technique for other data modalities such as images.