Cost-effective Distillation of Large Language Models (2023)

Abstract: Knowledge distillation (KD) involves training a small “student” model to replicate the strong performance of a high-capacity “teacher” model, enabling efficient deployment in resource-constrained settings. Top-performing methods tend to be task- or architecture-specific and lack generalizability. Several existing approaches require pretraining of the teacher on task-specific datasets, which can be costly for large and unstable for small datasets. Here we propose an approach for improving KD through a novel distillation loss agnostic to the task and model architecture. We successfully apply our method to the distillation of the BERT-base and achieve highly competitive results from the distilled student across a range of GLUE tasks, especially for tasks with smaller datasets.

## introduction an unfortunate problem affecting large language models, such as bert @xcite or gpt @xcite , is their high compute costs, as a consequence of their complex architectures and vast numbers of parameters. this is particularly apparent in initial (pre)training, but also impacts the cost of fine-tuning to specific tasks, and the practicality of their deployment on resource-constrained edge devices @xcite . knowledge distillation (kd; @xcite ) attempts to mitigate these concerns through learning a small "student" model to replicate the behaviour of a larger, unwieldy "teacher". the idea is that much of the performance of the teacher can be captured by the student, despite it having many fewer parameters, and thereby better portability. several distillation methods have been proposed for large language models, including distilbert @xcite , which distills the 12-layer bert transformer @xcite into a 6 layer student model with only a small loss in the performance on downstream tasks. broadly, existing kd approaches are either architecturespecific or agnostic. the former group includes @xcite and @xcite which incorporate a loss term to encourage matching hidden between teacher and student, and thus requiring aligned teacher and student architectures. approaches like @xcite , on the other hand are architecture-agnostic, treating the teacher model as a black box using only the logits from language modelling heads for distillation it into a smaller lm. there are numerous advantages to the architecture-agnostic approach: (1) it is possible to distill a teacher model into a different student architecture, e.g. @xcite distills the bert transformer into a simple single-layer bi-lstm; and (2) it frees the student to use different inference techniques, e.g., to better handle long sequences @xcite . while the training of large language models incurs substantial compute resources -for instance the training cost of gpt3 @xcite ) was estimated at @xmath00.11 an hour, which is considerably cheap compared to other approaches that include teacher fine-tuning. we showed that our method is particularly effective on small datasets, and competitive with other kd methods which are much more computationally intensive and tailored to the teacher. a possible reason could be since the fine-tuning of bert on small datasets like mrpc, stsb, or rte can be unstable @xcite , eliminating it makes the kd more robust and improves the results. all other methods such as tinybert @xcite or patientkd @xcite ) use fine-tuned teachers. distilbert @xcite does not use a fine-tuned teacher, but it is only limited to students with a hidden state of 784 due to the cosine loss it uses and lacks generalization across architectures.