Knowledge Distillation with Reptile Meta-Learning for Pretrained Language Model Compression (2022)

Abstract: The billions, and sometimes even trillions, of parameters involved in pre-trained language models significantly hamper their deployment in resource-constrained devices and real-time applications. Knowledge distillation (KD) can transfer knowledge from the original model (i.e., teacher) into a compact model (i.e., student) to achieve model compression. However, previous KD methods have usually frozen the teacher and applied its immutable output feature maps as soft labels to guide the student’s training. Moreover, the goal of the teacher is to achieve the best performance on downstream tasks rather than knowledge transfer. Such a fixed architecture may limit the teacher’s teaching and student’s learning abilities. Herein, a knowledge distillation method with reptile meta-learning is proposed to facilitate the transfer of knowledge from the teacher to the student. The teacher can continuously meta-learn the student’s learning objective to adjust its parameters for maximizing the student’s performance throughout the distillation process. In this way, the teacher learns to teach, produces more suitable soft labels, and transfers more appropriate knowledge to the student, resulting in improved performance. Unlike previous KD using meta-learning, the proposed method only needs to calculate the first-order derivatives to update the teacher, leading to lower computational cost but better convergence. Extensive experiments on the GLUE benchmark show the competitive performance achieved by the proposed method. For reproducibility, the code for this paper is available at: https://github.com/maxinge8698/ReptileDistil .

## introduction in recent years, pre-trained language models (plms) have brought natural language processing to a new era and achieved state-of-the-art performance in a variety of tasks. based on a multilayer transformer architecture @xcite , plms, such as bert @xcite , roberta @xcite , and xlnet @xcite , typically contain hundreds of millions of parameters, making them computationally expensive and inefficient. consequently, such large-scale models are difficult to be deployed on resource-constrained devices and real-time applications due to high computational complexity, huge storage requirements, and slow inference speed. knowledge distillation (kd) @xcite has been recommended to compress plms. based on the teacher-student architecture applied, it teaches a smaller student model to reproduce the behavior of a larger teacher model. in practice, the teacher produces output feature maps as soft labels, providing the student with more information than the ground truth of one-hot labels to learn. by minimizing the kullback-leibler (kl) divergence between the softened probability distributions of the teacher and student, the student is then trained as an equally-effective model without significant sacrifice in performance. however, previous studies @xcite @xcite @xcite have usually frozen the teacher during the kd process and used its immutable knowledge to teach the student. inevitably, this architecture has two major limitations: 1) the parameters of the teacher were fixed for kd. concretely, the teacher is frozen throughout the distillation process, and the student can only passively receive the fixed knowledge conveyed by the teacher, which ignores the capacity gap and incompatibility between the teacher and student; 2) the teacher is unaware of the existence of the student. specifically, the teacher model is usually fine-tuned to optimize its performance in downstream applications rather than to optimize the ability to distill the knowledge to the student. recent studies @xcite @xcite have suggested that up-dating the teacher together with the student during the kd process, rather than freezing it, can make the student learn better. to be specific, @xcite proposed a student-friendly teacher network for training the teacher and student branches to make the teacher aware of the student before distillation and obtain student-friendly knowledge, which is then transferred to the student via vanilla kd. @xcite presented meta pseudo labels where the student is trained based on the pseudo labels generated by the teacher and the teacher is trained based on the performance of the student on labeled data. @xcite employed a meta-learning strategy to explicitly optimize the teacher with the optimization objective of the student's performance on a hold-out training subset during the kd process, allowing the teacher to evolve continuously to adapt to the current state of the student and better transfer knowledge to it. the application of model-agnostic meta-learning (maml) @xcite in metadistil @xcite significantly improved the performance of the distilled student by calculating the second-order derivatives of the student's loss on the query set to obtain the update gradients of the teacher. however, limited by the maml algorithm, there are some concerns in metadistil: 1) a query set needs to be designed. metadistil split the query set from the training set at a ratio of 9:1, as a result, it missed one-tenth of the training data since this separate query set is not directly used to train the student model. nevertheless, this one-tenth of the data could lead to significant performance gain, especially when the training data is scarce; 2) the training process may be slow due to the nature of meta-learning. furthermore, maml involves the calculation of the second-order derivatives, so that more training time and computing resource are required. to address the above issues, a knowledge distillation method with reptile meta-learning @xcite is proposed to facilitate the transfer of knowledge from the teacher to the student, termed reptiledistil, where the update gradients of the teacher are approximated by calculating the difference between the parameters of the metalearner (i.e., teacher) and inner-learner (i.e., student), thereby reducing the computational burden by avoiding calculating the second-order derivatives, and mitigating the tedious operation of partitioning the query set and the performance loss that may be caused. to illustrate the effectiveness and practicality of the proposed reptiledistil, extensive experiments are conducted on the general language understanding evaluation (glue) benchmark (wang et al., 2018). empirical results demonstrate that the proposed method can yield meaningful improvements compared with other kd methods in distilling a 12-layer bert base model with 110m parameters into a 6-layer bert 6 model with 66m parameters and 1.94 times speedup. in summary, the contributions of this paper are as follows: the rest of this paper is organized as follows. section 2 reviews the related work. section 3 details the proposed reptiledistil framework. section 4 presents extensive experiments. section 5 provides experimental results and analysis. section 6 summarizes the paper. ## related work model compression is a potential approach to reduce the model size and improve computational efficiency. currently, the model compression techniques for plms can be divided into the following six categories: 1) model pruning @xcite @xcite , which removes redundant or less important parameters; 2) weight quantization @xcite , which uses fewer bits to represent the parameters; 3) knowledge distillation @xcite @xcite @xcite , which trains a smaller model that learns from the output feature maps of the original model; 4) module replacing @xcite , which replaces the modules of the original model with more compact substitutes; 5) matrix factorization @xcite , which reduces the parameters using methods such as cross-layer parameter sharing and factorized embedding parameterization; 6) early exit @xcite @xcite , which allows the model to exit early at an off-ramp instead of passing through the entire model. here, we briefly review state-of-the-art work on investigating knowledge distillation to compress plms, especially the popular bert model. dis-tilbert @xcite is performed at the pre-training stage, which distills output logits from a pre-trained bert base teacher into a 6-layer bert 6 student initialized by taking one layer out of every two layers of a pre-trained bert base . bert-pkd @xcite is performed at the fine-tuning stage, which distills output logits and hidden states from a fine-tuned bert base teacher into a 6-layer bert 6 student initialized with the first 6 layers of a pre-trained bert base . pd @xcite is performed at the pre-training stage, in which a randomly initialized 6-layer bert 6 student is first trained with a masked language modeling objective and then is distilled with the output logits from a pre-trained bert base teacher. tinybert @xcite first is performed at the pre-training stage, which distills embedding outputs, hidden states, and self-attention distributions from a pre-trained bert base teacher into a 4-layer tinybert 4 student with a hidden size of 312 and an intermediate size of 1200, or a 6-layer tinybert 6 student with a hidden size of 768 and an intermediate size of 3072 as same as bert 6 . these two models are then treated as the initialization of student models for further distillation. after that, at the fine-tuning stage, the output logits, embedding outputs, hidden states, and self-attention distributions from a fine-tuned bert base teacher are distilled into the aforementioned tinybert 4 student or tinybert 6 student on the augmented task-specific dataset. minilm v2 @xcite is performed at the pre-training stage, which distills self-attention relation from the last transformer layer of a pre-trained bert base teacher into a randomly initialized 6-layer bert 6 student. metadistil @xcite is performed at the fine-tuning stage, which distills output logits from a pre-trained bert base teacher into a 6-layer bert 6 student obtained from the aforementioned pd @xcite via a meta-learning strategy. ## conclusions in this paper, a knowledge distillation method with reptile meta-learning was proposed to compress pre-trained language models, which also utilizes the current performance of the student model during each step of the distillation process as feedback to optimize the teacher model, allowing the teacher to learn to teach and thus produce more appropriate soft labels to teach the student. however, unlike previous knowledge distillation methods applying meta-learning, the proposed method avoids calculating the second-order derivatives and achieves better and faster convergence, owing to the integration of reptile meta-learning and the proposed layer-wise parameter update strategies. extensive experiments demonstrated the competitive performance achieved by the proposed method. future work attempts to continue to explore such a dynamic and interactive teacher-student distillation architecture to improve the performance of the distilled model further.