Performance-Guided LLM Knowledge Distillation for Efficient Text Classification at Scale (2024)

Abstract: Large Language Models (LLMs) face significant challenges at inference time due to their high computational demands. To address this, we present Performance-Guided Knowledge Distillation (PGKD), a cost-effective and high-throughput solution for production text classification applications. PGKD utilizes teacher-student Knowledge Distillation to distill the knowledge of LLMs into smaller, task-specific models. PGKD establishes an active learning routine between the student model and the LLM; the LLM continuously generates new training data leveraging hard-negative mining, student model validation performance, and early-stopping protocols to inform the data generation. By employing a cyclical, performance-aware approach tailored for highly multi-class, sparsely annotated datasets prevalent in industrial text classification, PGKD effectively addresses training challenges and outperforms traditional BERT-base models and other knowledge distillation methods on several multi-class classification datasets. Additionally, cost and latency benchmarking reveals that models fine-tuned with PGKD are up to 130X faster and 25X less expensive than LLMs for inference on the same classification task. While PGKD is showcased for text classification tasks, its versatile framework can be extended to any LLM distillation task, including language generation, making it a powerful tool for optimizing performance across a wide range of AI applications.

## introduction successful research in artificial intelligence (ai) and large language models (llms) brought significant interest to the industrial applications of llms. llms mainly refer to transformerbased @xcite neural language models that contain tens to hundreds of billions of parameters, which are pre-trained on massive text * equal contribution. data, such as mistral 7b @xcite , llama @xcite , and @xcite . compared to smaller pre-trained language models (plms) like bert @xcite and gpt-2 @xcite , llms not only have significantly larger model sizes but also demonstrate stronger language understanding and generation capabilities. it has been demonstrated that llms can achieve considerable performance with limited task-specific annotated data across several domains, including natural language understanding, question answering, and code generation @xcite . in the text classification domain, llms like gpt-3 outperforms ad-hoc state-of-the-art (sota) models on several benchmarks by just using few-shot prompting techniques @xcite . while the utility of llms is evident, their application at scale is challenged by high inference latency and cost. llms are unsuitable for production environments where service-level agreements (slas) are strict and scalability is essential. furthermore, many natural language processing (nlp) tasks in the industry only require a subset of the llms capabilities, e.g., intent detection, a multi-class classification problem that can be solved by fast and inexpensive plm and does not require language generation capabilities. furthermore, utilizing llms for text classification restricts the nuanced output information derived from fine-tuning a plm. plms, when equipped with a classification head, produce output vectors with predicted probabilities that reflect the model's certainty about the classification results. however, this level of detail is not easily attainable from llms. llms, instead, generate plain text that requires parsing the model output and handling hallucinations, complicating production systems even further for simple tasks like text classification. considering the challenges presented by llms along with the effectiveness they bring, we propose a novel method called performance-guided knowledge distillation (pgkd). the proposed methodology aims to enhance the performance of smaller plms, which already meet inference sla and cost constraints, by using llms as teachers during training. ## related work knowledge distillation (kd), introduced by @xcite , is a promising technique to transfer the capabilities of complex and highmaintenance models to more compact and efficient student models. the core idea is to train a lean student model to mimic the soft probabilities generated by a more complex and costly teacher model. using llms for ground truth generation and as kd teacher has found success in multiple tasks @xcite @xcite . recent research @xcite has explored novel kd with open-source llms and loss functions explicitly tailored for teacherstudent distillation. the same work also summarizes the two commonly applied categories of kd: white-box kd @xcite , where the teacher parameters are available to use for model distillation, and black-box kd, where only the teacher predictions are accessible. black-box kd is less restrictive in terms of structural requirements for teacher models and student models, and can use closed-source teacher models (such as chatgpt api); additionally, it does not require the private deployment of teacher models and does not present challenges in custom loss function convergence. most recent black-box kd methods mainly focus on llm data generation or augmentation. @xcite proposed auggpt, a text data augmentation approach based on chatgpt that rephrases each sentence in the training samples into multiple conceptually similar but semantically different samples. zerogen @xcite is a flexible and efficient zero-shot learning method, providing insights on data-free, model-agnostic knowledge distillation. zeroshotdataaug @xcite investigates chatgpt-generated synthetic training data to augment low-resource datasets, outperforming existing data augmentation approaches, and explores methodologies for evaluating the quality of the generated augmented data. sungen @xcite optimizes synthetic data generation in black-box knowledge distillation by employing adversarial sampling and iterative refinement to create diverse, challenging data that aligns with the teacher model's knowledge boundaries. despite these advancements, the potential of llms in kd extends beyond mere data generation. recent studies have explored dynamic interactions between teacher and student models at training time for effective kd. @xcite proved the effectiveness of active learning for teacher-student distillation for llms and introduces the evokd framework. this framework is specifically tailored to enhance the capabilities of a student model within a few-shot learning scenario, where only a limited quantity of training data is available. research on text classification using kd has been focused on few-shot learning scenarios, such as 1-shot or 5-shot, where 1 to 5 training samples are provided to the base model, typically narrowing the task to binary classification. while this focus offers valuable insights for theoretical exploration, it does not adequately represent the complexity of text classification tasks encountered in industrial applications. in a typical industrial setting, classification challenges, such as intent detection, topic classification, and customer feedback analysis, involve a much larger number of categories. moreover, in these settings, annotating hundreds to a few thousand samples is often feasible and economically viable, making the 1-shot or 5-shot evaluation unrealistic given the availability of annotated data. ## methods pgkd is an iterative kd algorithm that aims to enhance the application of kd in highly multi-class, sparsely annotated datasets common in industrial environments. the primary motivation for this work is to leverage active learning to strengthen the connection between the student and teacher models in the distillation process. pgkd gives the teacher model direct and continued visibility into the learning status of the student model also making it generically extensible to wide variety of learning tasks. pgkd addresses limitations of recent works in black-box kd, in particular: (1) evokd only shows results on binary-classification benchmark datasets, a setting that is infrequent in practical machine learning applications; (2) ze-rogen only relies on the latent knowledge of the teacher and is not aware of the student's status; the teacher model in evokd is also not provided with an overview of the latest student model performance on all the available classes and only shows the llm a few misclassified examples; while this approach is useful in the context of binary text classification, it can be limiting in a highly multi-class context where many classes might not be represented in the sample provided to the teacher. (3) in sungen, distillation focuses on the capability of the teacher model to generate high entropy data; evokd too focusses on correctly classified and misclassified samples but ignores emphasizing the information that is contained in the "hard negative samples", i.e., incorrectly classified samples which the student model is confident about. (4) previous works do not provide a clear termination criterion for the iterative algorithm but treat the number of epochs for the knowledge distillation process as a hyperparameter that needs to be tuned for each dataset. the proposed pgkd overcomes the limitation above by leveraging the following techniques. the pgkd methodology is detailed in algorithm 1 and a schema of the same is reported in figure 1 and figure 2 . the algorithm starts by initializing and training a baseline model on an initial dataset d 0 composed of a thousand annotated samples, divided into 80% training and 20% validation. the base model obtained after this training is denoted as model 0 . pgkd iteratively refines model 0 performance; during each iteration, the pgkd algorithm identifies correctly classified d i correct and misclassified examples d i incorrect , as well as hard negatives d i hard_negatives . pgkd then computes the model val_results on the validation set d val and leverages an llm to generate new training data d i+1 tailored to these findings according to p gkd_prompt. p gkd_prompt includes dataset classification taxonomy for the specific task and a limited number of few-shot samples from the train set, prompt is reported in appendix a. the pgkd process runs for num_kd_steps and if the validation loss does not improve consecutively beyond the patience_limit, the algorithm terminates early to prevent overfitting. at the end of the pgkd process, the best model on the validation set, model best is returned. ## datasets and experiments our experiments focus on four multi-class classification datasets: ag-news @xcite , yahoo answers @xcite , huffington post (misra, 2018), and amzn reviews (kashnitsky, 2018), described in if @xmath0: @xcite ; however, further llm prompt exploration is beyond the scope of this work. while pgkd is expected to be maximally useful when the number of annotated data is limited, it is interesting to evaluate performance gains obtained by pgkd as the number of training samples increases. figure 3 presents the results for the amzn reviews dataset obtained at a varying number of training dataset sizes, with similar trends observed across all datasets; full results are documented in appendix b. the pgkd approach consistently surpassed the performance of the bert-base model. notably, as the number of training samples from the original dataset increases, the performance differential between pgkd and the original model decreases. this phenomenon can be attributed to the diminishing returns of new data generation by the llm as the volume of original dataset samples grows. interestingly, while pgkd performance improvement reduces with the increasing number of training samples, we can observe that the pgkd provides the best results across all dataset considered for all training sample sizes. pgkd demonstrates to not degrade model in any of the experiments conducted. ## conclusion and future work deploying llms in real-world text classification applications poses significant challenges due to high inference costs and latency. this research introduces pgkd, a novel methodology designed to effectively distill llm knowledge into faster, more efficient models for multi-class text classification. this study proves the substantial performance improvements of pgkd distillation compared to regular fine-tuning of a bert-base plm, while maintaining limited inference latency and costs. comparative analyses with existing knowledge distillation and augmentation strategies further underscore pgkd's practical performance enhancements. the proposed ablation study reveals the significant contributions of pgkd components, such as gradual performance checks on validation reports and hard negative mining. cost and latency benchmarks provide compelling evidence of pgkd's efficiency. compared to zero-shot llm costs, bert-base with pgkd has proven to be significantly more cost-effective and up to 130x faster for a broad spectrum of multi-class classification tasks. future research will investigate the impact of different teacher llms on the distillation process and explore the influence of the student model size on pgkd effectiveness. future research will also focus on exploring more advanced prompting techniques to optimize pgkd performance. while pgkd is showcased for text classification tasks, its versatile framework can be extended to any llm distillation task, including language generation, making it a powerful tool for optimizing performance across a wide range of ai applications.