R e A ug KD : Retrieval-Augmented Knowledge Distillation For Pre-trained Language Models (2023)

Abstract: Knowledge Distillation (KD) is one of the most effective approaches to deploying large-scale pre-trained language models in low-latency environments by transferring the knowledge contained in the large-scale models to smaller student models. Prior KD approaches use the soft labels and intermediate activations generated by the teacher to transfer knowledge to the student model parameters alone. In this paper, we show that having access to non-parametric memory in the form of a knowledge base with the teacher’s soft labels and predictions can further improve student generalization. To enable the student to retrieve from the knowledge base effectively, we propose a new framework and loss function that preserves the semantic similarities of teacher and student training examples. We show through extensive experiments that our retrieval mechanism can achieve state-of-the-art performance for task-specific knowledge distillation on the GLUE benchmark.

## introduction large pre-trained language models, such as bert @xcite , roberta @xcite and electra @xcite have achieved significant success on several different nlp tasks @xcite with finetuning. however, these models usually contain millions and billions of parameters, preventing their execution on resource-restricted devices. to deploy these models, knowledge distillation (kd) is an effective compression technique to derive a smaller student model from a larger teacher model by transferring the knowledge embedded in the teacher's network. previous kd methods typically store knowledge in the student's parameters and train the student by minimizing divergence between the student's and teacher's output prediction and intermediate activation distributions @xcite . however, the student's parametric memory is often limited and cannot be quickly expanded or revised. moreover, after training, the teacher model's soft labels and activations, which contain essential task-specific knowledge, are not utilized by the student at inference time. to address the issues mentioned above, we propose the retrieval-augmented knowledge distillation (reaugkd) framework. reaugkd introduces a non-parametric external memory in addition to the implicit parametric memory of the model and uses knn retrieval to retrieve from this memory. the key intuition of reaugkd is to enhance the effective capacity of the student by using an external memory derived from relevant task-specific knowledge of the teacher. while this external memory could include any task-specific knowledge, in this work, it is composed of the soft labels and embeddings generated by the teacher model. our framework consists of an inference phase and a training phase. in the inference phase, we aggregate the soft labels of those teacher embeddings in our memory that are most similar to the student embedding. we demonstrate the efficacy of our framework by achieving state-of-the-art results on the glue benchmark @xcite with less than 3% latency overhead over the baseline without retrieval augmentation. reaugkd also comprises a training phase, where we train the student to retrieve from the external memory effectively. we train with a novel relational kd loss that minimizes the divergence between teacher-teacher and teacher-student embedding distributions. we not only observe that training with this loss is necessary to align the student and teacher embedding spaces for retrieval but also that this loss improves student generalization even in the absence of retrieval augmentation. this suggests that incorporating the ability to retrieve information can significantly enhance generalization during the process of knowledge distillation. in summary, our contributions include • we propose reaugkd, a novel framework for knowledge distillation that introduces a nonparametric memory to increase the effective student size. we show that retrieving from a memory composed of training set teacher predictions at inference time can significantly improve generalization on the glue tasks. ## conclusion in this paper, we present reaugkd, a knowledge distillation framework with a retrieval mechanism that shows state-of-the-art performance on the glue benchmark. in the future, we plan to expand the knowledge base with more information from the teacher and extend it to additional tasks.