Marginal Utility Diminishes: Exploring the Minimum Knowledge for BERT Knowledge Distillation (2021)

Abstract: Recently, knowledge distillation (KD) has shown great success in BERT compression. Instead of only learning from the teacher’s soft label as in conventional KD, researchers find that the rich information contained in the hidden layers of BERT is conducive to the student’s performance. To better exploit the hidden knowledge, a common practice is to force the student to deeply mimic the teacher’s hidden states of all the tokens in a layer-wise manner. In this paper, however, we observe that although distilling the teacher’s hidden state knowledge (HSK) is helpful, the performance gain (marginal utility) diminishes quickly as more HSK is distilled. To understand this effect, we conduct a series of analysis. Specifically, we divide the HSK of BERT into three dimensions, namely depth, length and width. We first investigate a variety of strategies to extract crucial knowledge for each single dimension and then jointly compress the three dimensions. In this way, we show that 1) the student’s performance can be improved by extracting and distilling the crucial HSK, and 2) using a tiny fraction of HSK can achieve the same performance as extensive HSK distillation. Based on the second finding, we further propose an efficient KD paradigm to compress BERT, which does not require loading the teacher during the training of student. For two kinds of student models and computing devices, the proposed KD paradigm gives rise to training speedup of 2.7x 3.4x.

## introduction since the launch of bert @xcite , pre-trained language models (plms) have been advancing the state-of-the arts (sotas) in a wide range of nlp tasks. at the same time, the growing * work was done when yuanxin liu was an intern at pattern recognition center, wechat ai, tencent inc, china. figure 1 : the acc variation of rosita @xcite and tinybert @xcite on qnli with the increase of hsk. in contrast to the previous work that attempts to increase the amount of hsk, in this paper we explore towards the opposite direction to "compress" hsk. we make the observation that although distilling hsk is helpful, the marginal utility diminishes quickly as the amount of hsk increases. to understand this effect, we conduct a series of analysis by compressing the hsk from three dimensions, namely depth, length and width (see section 2.3 for detailed description). we first compress each single dimension and compare a variety of strategies to extract crucial knowledge. then, we jointly compress the three dimensions using a set of compression configurations, which specify the amount of hsk assigned to each dimension. figure 1 shows the results on qnli dataset. we can find that 1) perceivable performance improvement can be obtained by extracting and distilling the crucial hsk, and 2) with only a tiny fraction of hsk the students can achieve the same performance as extensive hsk distillation. based on the second finding, we further propose an efficient paradigm to distill hsk. concretely, we run bert over the training set to obtain and store a subset of hsk. this can be done on cloud devices with sufficient computational capability. given a target device with limited resource, we can compress bert and select the amount of hsk accordingly. then, the compressed model can perform kd on either the cloud or directly on the target device using the selected hsk and the original training data, dispensing with the need to load the teacher model. in summary, our maojor contributions are: ## single-dimension knowledge compression researches on model pruning have shown that the structural units in a model are of different levels of importance, and the unimportant ones can be dropped without affecting the performance. in this section, we investigate whether the same law holds for hsk compression in kd. we study the three dimensions separately and compare a variety of strategies to extract the crucial knowledge. when a certain dimension is compressed, the other two dimensions are kept to full scale. ## three-dimension joint knowledge compression with the findings in single-dimension compression, we are now at a position to investigate joint hsk compression from the three dimensions. ## improving training efficiency existing bert compression methods mostly focus on improving the inference efficiency. however, the teacher model is used to extract features throughout the training process, which suggests that the training efficiency still has room for improvement. as shown in figure 10 , the compressed models achieve considerable inference speedup, while the increase in training speed is relatively small. moreover, for students with different sizes or architectures, the teacher should be deployed every time when training a new student. intuitively, we can run the teacher once and reuse the features for all the students. in this way, we do not need to load the teacher model while training the student, and thereby increasing the training speed. we refer despite the training speedup, however, loading and storing hsk increases the memory consumption. the full set of hsk can take up a large amount of space, especially for the pre-trained language models like bert. fortunately, our findings in the previous sections suggest that the student only requires a tiny fraction of hsk. we explore a variety of strategies to determine feature importance for each single dimension. this is related to a line of studies called the attribution methods, which attempt to attribute a neural network's prediction to the input features. the attention weights have also been investigated as an attribution method. however, prior work @xcite @xcite finds that attention weights usually fail to correlate well with their contributions to the final prediction. this echoes with our finding that the original att strategy performs poorly in length compression. however, the attention weights may play different roles in attribution and hsk distillation. whether the findings in attribution are transferable to hsk distillation is still a problem that needs further investigation. ## conclusions and future work in this paper, we investigate the compression of hsk in bert kd. we divide the hsk of bert into three dimensions and explore a range of compression strategies for each single dimension. on this basis, we jointly compress the three dimensions and find that, with a tiny fraction of hsk, the student can achieve the same or even better performance as distilling the full-scale knowledge. based on this finding, we propose a new paradigm to improve the training efficiency in bert kd, which does not require loading the teacher model during training. the experiments show that the training speed can be increased by 2.7× ∼ 3.4× for two kinds of student models and two types of cpu and gpu devices. most of the compression strategies investigated in this study are heuristic, which still have room for improvement. therefore, a future direction of our work could be designing more advanced algorithm to search for the most useful hsk in bert kd. additionally, since hsk distillation in the pre-training stage is orders of magnitude time-consuming than task-specific distillation, the marginal utility diminishing effect in pre-training distillation is also a problem worth studying. yoshikiyo kato, taku kudo, hideto kazawa, keith stevens, george kurian, nishant patil, wei wang, cliff young, jason smith, jason riesa, alex rudnick, oriol vinyals, greg corrado, macduff hughes, and jeffrey dean. 2016. google's neural machine translation system: bridging the gap between human and machine translation. corr, abs/1609.08144.