ReAugKD: Retrieval-Augmented Knowledge Distillation For Pre-trained Language Models
Year: 2023

Abstract:
Knowledge Distillation (KD) (Hinton et al., 2015) is one of the most effective approaches for deploying large-scale pre-trained language models in low-latency environments by transferring the knowledge contained in the largescale models to smaller student models. Previous 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 enhance student capacity and improve generalization. To enable the student to retrieve from the knowledge base effectively, we propose a new Retrieval-augmented KD framework with a loss function that aligns the relational knowledge in teacher and student embedding spaces. We show through extensive experiments that our retrieval mechanism can achieve state-of-the-art performance for taskspecific knowledge distillation on the GLUE benchmark (Wang et al., 2018a).

heading: 1 Introduction
text: Large pre-trained language models, such as BERT (Devlin et al., 2018), RoBERTa (Liu et al., 2019) and Electra (Clark et al., 2020) have achieved significant success on several different NLP tasks (Ding et al., 2019; Wang et al., 2018a) 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 (Park et al., 2019; Zhang et al., 2018). 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 (Wang et al., 2018a) 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
1128
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. • To effectively retrieve from the non-parametric memory, we introduce a novel loss function that transfers the relational knowledge between teacherteacher embedding and teacher-student embedding distribution. This loss function improves student generalization even in the absence of retrieval augmentation at inference time. • We study the accuracy and latency cost with the number of neighbors (k) retrieved in ReAugKD. ReAugKD with approximate kNN introduces a small overhead of <3% latency increase.

heading: 2 Related Work
text: Knowledge distillation KD can be broadly classified into task-specific KD, where the student model will be used for the same task as the teacher model (Mirzadeh et al., 2020; Jin et al., 2019; Zhang et al., 2018; Sun et al., 2019) and task-agnostic KD where the student may be used for a different task, after finetuning on the new task (Jiao et al., 2019; Sun et al., 2020; Sanh et al., 2019; Wang et al., 2020; Zhang et al., 2018; Xu et al., 2019). In this work, we show that ReAugKD can be applied to enhance task-specific distillation as well as when finetuning task-agnostic distilled models. Closest to our work is RKD (Park et al., 2019) that introduces a loss to transfer relational knowledge between teacherteacher embedding and student-student embedding distributions. Our work differs in that we transfer relational knowledge between teacher-teacher embedding and teacher-student embedding distribution to enhance the student model’s ability to retrieve from the external memory. MetaDistil (Zhou et al., 2022) is a strong task-specific distillation baseline that employs meta-learning to better transfer knowledge to the student. Unlike MetaDistill, we show that ReAugKD can significantly improve the student model’s generalization without retraining the whole teacher with meta-learning. Retrieval-augmented language models There has been growing interest in retrieval-augmented methods for Knowledge-Intensive generative NLP
Tasks, such as text generation and question answering (Weston et al., 2018; Lewis et al., 2020; Guu et al., 2020; Lin et al., 2022), where querying training examples during inference significantly improves likelihood. Closest to our work is BERTkNN (Kassner and Schütze, 2020) which combines BERT with a kNN search over a large datastore of an embedded text collection, to improve clozestyle QA. In our work, we apply retrieval augmentation to enhance the capacity of student models during KD, and show improvement even on nonknowledge intensive tasks like GLUE.

heading: 3 Methodology

heading: 3.1 Training Phase
text: Our framework consists of two main phases, the training phase and the inference phase. The training phase has two steps. In the first step, we prepare the teacher model for KD by adding a linear projection head L on the top of the teacher model encoder that has been finetuned for a specific downstream task. The input dimension of this projection head is the embedding dimension of the teacher. The output dimension is the embedding dimension of the student. We then freeze the other parameters of the teacher model and finetune the parameters in L with supervised contrastive loss (Khosla et al., 2020). This step a) reduces the dimension of the teacher’s embeddings, to the student model dimension for retrieval, and b) uses supervised contrastive loss to derive a kNN classifier for BERT that is robust to natural corruptions, and hyperparameter settings (Li et al., 2021). Fine-tuning L also greatly reduces the computational cost compared to retraining the whole teacher model (Zhou et al., 2022).
In the second step, we perform KD by generating the teacher embeddings with L and teacher soft labels using the original teacher’s classifier head for a batch of data. Then, we use the loss function we proposed in Section 3 to train our student model.

heading: 3.2 Loss function
text: We present some mathematical notations to introduce our loss function. Given a batch of data {di}, i = 1, 2, · · · , N , where N is the batch size, we denote the embedding generated by the teacher’s projection head as zi and the soft labels generated by the teacher’s classifier as ȳi. Similarly, we adopt xi, yi to denote the student’s embeddings and predictions. Then we construct a probability distribution qi,j over each teacher’s embeddings zj
to capture the similarity with respect to an anchor point zi,
qi,j = exp (zi · zj)/τ∑N k=1 exp (zi · zk)/τ , (1)
where the τ stands for temperature. Note that∑N j=1 qi,j = 1. qi,j reflects the cosine distance relational knowledge among different embeddings generated by the teacher model in the batch. If zj is closer to zi, cosine distance, qi,j will be larger. Similarly, given a student’s embedding xi as an anchor point, we formulate another probability distribution q̄i,j over each teacher’s embeddings zj of the data in the batch.
q̄i,j = exp (xi · zj)/τ∑N k=1 exp (xi · zk)/τ . (2)
The q̄i,j reflects the cosine distance relationship between different embeddings generated by the teacher model and the student’s embedding. Our loss function aims to minimize the divergence of these two distributions q̄i,j and qi,j since the teacher model is a strong kNN classifier after finetuning with supervised contrastive loss function in the first step of our training. In the ideal case, given a student’s embedding xi, the student retriever should retrieve the same set of embeddings as the corresponding teacher’s embedding zi. We adopt KL divergence to measure that divergence. In addition, we adopt the commonly-used cross-entropy loss to calculate the divergence between the student’s prediction yi and the teacher’s prediction ȳi.
Our loss function can be formulated as
CE(yi, ȳi) + αKL(qi,j , q̄i,j), (3)
where CE is the cross entropy loss and KL is KLdivergence. α is the hyperparameter controlling the trade-off between the two losses.

heading: 3.3 Inference Phase
text: After training, we construct a knowledge base (KB) comprising of projected teacher embeddings and predictions. Given new data di at inference time, we obtain (xi, yi) using the student model. and use the HNSW algorithm (Malkov and Yashunin, 2018) to derive the K nearest teacher’s embeddings and their corresponding soft labels {(zk, ȳk)}i=1,2,··· ,K from the KB. Then we compute the weighted average of these soft labels Avg({ȳ})i based on q̄i,k
Avg({y})i = K∑
k=1
q̄i,k∑K k=1 q̄i,k ȳk
We derive a new prediction ȳ′i for di with Avg({ȳ})i.
ȳ′i = βȳi + (1− β)Avg({ȳ})i,
β is the hyperparameter controlling the trade-off between the two predictions.

heading: 4 Experimental Results
text: We apply our method to distill BERT-Base (Devlin et al., 2018) into a 6-layer BERT with a hidden size of 768. We evaluate our proposed approach, ReAugKD, on the GLUE benchmark (Wang et al., 2018a). These datasets can be broadly divided into three families of problems: single-set tasks that include linguistic acceptability (CoLA) and sentiment analysis (SST-2), similarity, and paraphrasing tasks (MRPC and QQP); inference tasks
that include Natural Language Inference (MNLI and RTE); and Question Answering (QNLI). We compare our method with vanilla KD (Hinton et al., 2015), TAKD (Mirzadeh et al., 2020), RCO (Jin et al., 2019), RKD (Park et al., 2019), DML (Zhang et al., 2018), PKD (Sun et al., 2019) ProKT (Shi et al., 2020), SFTN (Park et al., 2021) and MetaDistil (Zhou et al., 2022). Following similar setting as MetaDistill, we perform a grid search over the sets of the weight of KD loss from {0.9, 0.99}, the predictions weight β from {0, 0.1, ... 1} and the top-k from 1 to 20. We set the student learning rate to 2e-5 and the batch size to 64.
Experimental Results on GLUE We report the experimental results on the development set of the six GLUE tasks in Table 1. Notably, our method achieves start-of-the-art results on five out of the six datasets with an average improvement of 0.34% over the previous best KD method MetaDistil (Zhou et al., 2022). Although MetaDistil achieves slightly better performance on the MRPC dataset, our method has the advantage of not needing to conduct meta-learning on the whole large teacher model, which significantly increases extra training cost in terms of time and memory (Zhou et al., 2022). In addition, we also observe a performance gain of 0.37% with the retrieval component of ReAugKD as compared to ReAugKD without retrieval which verifies the benefit of retrieval augmentation in our approach. Even without the retrieval process, the student model trained by our
designed loss can still achieve comparable performance to MetaDistill on most datasets. Since our loss is designed to improve the student retrieval function, this demonstrates the importance of retrieval capability in KD.
Number of Neighbors Retrieved (k) To understand the time overhead of retrieval on the student model’s inference time, we investigate the performance and additional time overhead of the retrieval process while varying the number of neighbors retrieved (k) in Table 2. From the results, it is clear that retrieval improves the student model performance with an additional time overhead of less than 3% of the original inference time. The retrieval process is conducted only on CPU, and does not take up GPU resources during training.

heading: 5 Conclusion
text: 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.
Limitations Our method relies on having access to teacher embeddings and prediction which may not always be possible in a black-box distillation setting. Retrieval augmentation also requires maintaining a knowledge base that is memory intensive. The cost of the retrieval process is dependent on the size of the training corpus, which can be a limitation when dealing with very large training datasets. Conducting dataset distillation (Wang et al., 2018b) on the training corpus to further reduce memory cost and retrieval time is an important future step for our framework.
Acknowledgments This work was done when Jianyi Zhang was an intern at Amazon Search. In addition, Jianyi Zhang and Yiran Chen disclose support from grants CNS-2112562, IIS-2140247, and CNS-1822085. We thank Yuchen Bian for the valuable discussion and thank all reviewers for their valuable comments.

heading: A Appendix
text: A.1 ReAugKD with task-agnostic distillation
Previous results have demonstrated the effectiveness of our method for task-specific distillation. Our method can further improve the finetuned performance of task-agnostic distilled models. We adopt RoBERTa-large as the teacher model and the MiniLMv2 as the student model to verify the effectiveness of our method. Our method can achieve around 2% improvement in performance.
A.2 Details about training teacher model’s projection head We adopt the Lsupout version of the loss function in (Khosla et al., 2020) to finetune the parameters of the projection head, which is
Lsupout = − N∑
i=1
1
N
∑
j∈P (i) log exp (zi · zj) /τ∑N k=1 exp (zi · zk) /τ . (4)
Here, there are N data samples di in the batch and we denote the embedding generated by the teacher’s projection head for the i-th data di as zi. P (i) here represents the set of all the positive data samples for data di. The data samples from the same class are considered as positive pairs and the data samples from different classes are considered as negative pairs. Regarding the use of data augmentation in training the projection head, we chose not to adopt data augmentation as we found that using supervised contrastive loss without data augmentation was sufficient to achieve results comparable to the cross-entropy loss used in supervised learning. We use the AdamW optimizer with a learning rate of 0.00002. The batch size was set to 512, and the temperature for the supervised contrastive loss (SCL) was set to 0.07. We trained the model 3 epochs.
ACL 2023 Responsible NLP Checklist

heading: A For every submission:
text: 3 A1. Did you describe the limitations of your work?
It lies before the reference.
7 A2. Did you discuss any potential risks of your work? We think our work will not have any potential risk.
3 A3. Do the abstract and introduction summarize the paper’s main claims? Section 1
7 A4. Have you used AI writing assistants when working on this paper? Left blank.
B 7 Did you use or create scientific artifacts? Left blank.
B1. Did you cite the creators of artifacts you used? No response.
B2. Did you discuss the license or terms for use and / or distribution of any artifacts? No response.
B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? No response.
B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? No response.
B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and linguistic phenomena, demographic groups represented, etc.? No response.
B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. No response.
C 3 Did you run computational experiments? Section 3
3 C1. Did you report the number of parameters in the models used, the total computational budget (e.g., GPU hours), and computing infrastructure used? Section 3
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.
3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found hyperparameter values? Section 3
3 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? Section 3
7 C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? The packages we used are confidential due to our company’s policy
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? No response.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? No response.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? No response.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? No response.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? No response.

heading: Future Work
text: In the future, we plan to expand the knowledge base with more information from the teacher and extend it to additional tasks. Limitations Our method relies on having access to teacher embeddings and prediction which may not always be possible in a black-box distillation setting. Retrieval augmentation also requires maintaining a knowledge base that is memory intensive. The cost of the retrieval process is dependent on the size of the training corpus, which can be a limitation when dealing with very large training datasets. Conducting dataset distillation (Wang et al. , 2018b) on the training corpus to further reduce memory cost and retrieval time is an important future step for our framework. Acknowledgments This work was done when Jianyi Zhang was an intern at Amazon Search. In addition, Jianyi Zhang and Yiran Chen disclose support from grants CNS-2112562, IIS-2140247, and CNS-1822085