Are Intermediate Layers and Labels Really Necessary? A General Language Model Distillation Method (2023)

Abstract: The large scale of pre-trained language models poses a challenge for their deployment on various devices, with a growing emphasis on methods to compress these models, particularly knowledge distillation. However, current knowledge distillation methods rely on the model’s intermediate layer features and the golden labels (also called hard labels), which usually require aligned model architecture and enough labeled data respectively. Moreover, the parameters of vocabulary are usually neglected in existing methods. To address these problems, we propose a general language model distillation (GLMD) method that performs two-stage word prediction distillation and vocabulary compression, which is simple and surprisingly shows extremely strong performance. Specifically, GLMD supports more general application scenarios by eliminating the constraints of dimension and structure between models and the need for labeled datasets through the absence of intermediate layers and golden labels. Meanwhile, based on the long-tailed distribution of word frequencies in the data, GLMD designs a strategy of vocabulary compression through decreasing vocabulary size instead of dimensionality. Experimental results show that our method outperforms 25 state-of-the-art methods on the SuperGLUE benchmark, achieving an average score that surpasses the best method by 3%.

## introduction the exponential increase in the scale of pre-trained language models has impeded their deployment on a wider range of devices. to mitigate the inference cost of large-scale pre-trained language models, researchers have increasingly focused on model compression methods, aiming to compress a large model into a small one with as little performance loss as possible @xcite . while model compression can yield very small models, maintaining performance without degradation is still a challenging task, particularly when the large and small models have a significant discrepancy in parameter size @xcite . there are various methods of model compression @xcite , including network pruning @xcite , quantization @xcite , neural architecture search @xcite , parameter sharing @xcite , matrix decomposition @xcite , and knowledge distillation @xcite etc. currently, knowledge distillation is an important research direction, which allows for the transfer of knowledge from a large model (the teacher) to a small one (the student). there are two main optimization objectives of the earliest knowledge distillation methods @xcite : increasing the similarity between the student's prediction probabilities for the task and those of the teacher (soft targets); increasing the similarity between the student's predictions and the golden labels (hard targets). when the knowledge distillation method is applied to language models, there are typically two directions for improvement: leveraging the intermediate layer features of the teacher model, such as hidden states and attention, to obtain additional hidden state knowledge @xcite ; and refining the two objectives (soft targets and hard targets) and weights of objectives @xcite . as shown in to address these problems, we propose a general language model distillation (glmd) method that performs two-stage (pre-training and task-specific stages) word prediction distillation and vocabulary compression. specifically, glmd distills the model using only the language modeling word prediction logits during the pre-training stage, which is similar to the soft labels used in general methods. the key to this stage is that we distill both masked and unmasked tokens. in the task-specific stage (fine-tuning), glmd distills both the language modeling word prediction logits and the soft labels. the language modeling word prediction logits is crucial in this stage, making the distillation more consistent between the pre-training and task-specific stages. in these two stages, glmd eliminates the need for complicated intermediate layers and golden labels and does not require the selection of intermediate layers or labeled dataset. meanwhile, glmd uses the teacher vocabulary to map low-frequency words to the most similar highfrequency words, further compressing the model with almost no performance loss. in summary, our major contributions are: we demonstrate through analysis that glmd allows models to autonomously learn intermediate layer features that are similar to those of the teacher. ## related work language model distillation since the introduction of knowledge distillation to pre-trained language models by pkd @xcite , an increasing number of researchers have recognized the importance of knowledge distillation. during the early stage of the research, pd @xcite employed simple baseline (soft targets) distillation for language models, resulting in a relatively limited transfer of knowledge for the model. subsequent research had primarily focused on the use of intermediate layer features in language models @xcite , including distillation of models during pre-training stage @xcite , task-specific stage @xcite , and two-stage @xcite approaches. given the typically large amount of intermediate layer features, some work had utilized features from only a single intermediate layer @xcite , while other work had examined methods for reducing the scale of features @xcite . recent work has explored ways to utilize better intermediate layer features, for example, codir @xcite and lrc-bert @xcite utilized cross-sample feature relationships through contrastive learning; alp-kd @xcite and universal-kd @xcite ) combined all intermediate layer features through attention mechanisms; meta-kd @xcite and hrkd @xcite used meta-learning to assign appropriate weights to intermediate layer features; rail-kd (haidar et al., 2022) randomly selected different intermediate layers for distillation; ckd @xcite and mgskd @xcite used some variety of similarity calculation methods for intermediate layer features; diito @xcite allowed student models to learn counterfactual outputs by swapping intermediate layer features between different samples. however, the use of intermediate layer features has additional limitations, such as requiring the same model structure @xcite for both teacher and student, or requiring linear transformations @xcite to ensure consistency in dimensions between teacher and student. there were also methods that only used soft and hard targets, for example, annealing-kd @xcite and continuation-kd @xcite gradually increased the weight of soft targets through simulated annealing; rw-kd @xcite adjusted the weight of soft and hard targets through meta-learning and a dev set; metadistil @xcite allowed the teacher to learn how to output better soft labels through meta-learning and a quiz set. these approaches relied on hard labels and may have even required additional datasets for partitioning. additionally, there had been approaches that distilled multiple teachers @xcite @xcite or teacher assistants @xcite at the same time, but they still relied on intermediate layer features or hard labels. in comparison, glmd can achieve the strongest performance in a more broadly applicable context without intermediate layer features or hard labels. vocabulary compression vocabulary compression refers to reducing the parameter size of the vocabulary in a language model. in the knowledge distillation of language models, reducing the parameter size of the model is mainly used to reduce the number of model layers or dimensions @xcite . mobilebert @xcite and albert @xcite independently reduced the dimensions of the vocabulary to achieve vocabulary compression. mo-bilebert needed to restore the dimension of the vocabulary in the calculation of the pre-training loss due to the requirement to ensure consistency between the vocabulary dimension and the model output dimension. on the other hand, albert used a linear layer to alter the output dimension of the model. however, these vocabulary compression methods only reduced the dimensionality and ultimately required dimensionality restoration. in contrast, our vocabulary compression method reduces the number of words through mapping, further compressing the model with almost no impact on performance. ## preliminaries in this section, we introduce the objective function for knowledge distillation of language models, and formalize the language modeling word prediction logits. ## method in this section, we propose a general language model distillation (glmd) method with two-stage word prediction distillation and vocabulary compression. figure 1 shows the overview framework of glmd, which implements a vocabulary compression strategy while performing a two-stage transformers embeddigs-layer reuse transformers embeddigs-layer pre-trained teacher student model transformers embeddigs-layer transformers embeddigs-layer student model fine-tuned teacher word prediction distillation process. we next provide a detailed description of these two components. ## experiments in this section, we demonstrate the effectiveness of the glmd on models with different parameter scales (110m, 340m, and 10b) and analyze the role of different components and why they are effective. all experiments were conducted on 40 nvidia a100 gpus and completed within 4 months, utilizing the pytorch framework. ## analysis in this section, we analyze the reasons behind the work of glmd and the impact of hyperparameters on performance in glmd. ## conclusions in this paper, we introduce a general language model distillation method called glmd. glmd has two main advantages: improving distillation performance without relying on intermediate layer features and hard labels and reducing vocabulary parameters without reducing dimensions. we also had two important findings: distillation of intermediate layer features is unnecessary, and a vocabulary compression strategy that reduces the number of tokens is feasible and can be combined with a method that reduces dimensions. in the future, we plan to explore model distillation on a 100b-scale and apply it to more real-world scenarios.