Distilling Linguistic Context for Language Model Compression (2021)

Abstract: A computationally expensive and memory intensive neural network lies behind the recent success of language representation learning. Knowledge distillation, a major technique for deploying such a vast language model in resource-scarce environments, transfers the knowledge on individual word representations learned without restrictions. In this paper, inspired by the recent observations that language representations are relatively positioned and have more semantic knowledge as a whole, we present a new knowledge distillation objective for language representation learning that transfers the contextual knowledge via two types of relationships across representations: Word Relation and Layer Transforming Relation. Unlike other recent distillation techniques for the language models, our contextual distillation does not have any restrictions on architectural changes between teacher and student. We validate the effectiveness of our method on challenging benchmarks of language understanding tasks, not only in architectures of various sizes but also in combination with DynaBERT, the recently proposed adaptive size pruning method.

## introduction since the transformer, a simple architecture based on attention mechanism, succeeded in machine translation tasks, transformer-based models have become a new state of the arts that takes over more complex structures based on recurrent or convolution networks on various language tasks, e.g., language understanding and question answering, etc @xcite @xcite @xcite . however, in exchange for high performance, these models suffer from a major drawback: tremendous computational and memory costs. in particular, it is not possible to deploy such large models on platforms with limited resources such as mobile and wearable devices, and it is an urgent research topic with impact to keep up with the performance of the latest models from a small-size network. as the main method for this purpose, knowledge distillation (kd) transfers knowledge from the large and well-performing network (teacher) to a smaller network (student). there have been some efforts that distill transformer-based models into compact networks @xcite @xcite @xcite . however, they all build on the idea that each word representation is independent, ignoring relationships between words that could be more informative than individual representations. in this paper, we pay attention to the fact that word representations from language models are very structured and capture certain types of semantic and syntactic relationships. -word2vec @xcite and glove @xcite demonstrated that trained embedding of words contains the linguistic patterns as linear relationships between word vectors. recently, @xcite found out that the distance between words contains the information of the dependency parse tree. many other studies also suggested the evidence that contextual word representations @xcite and attention matrices @xcite contain important relations between words. moreover, @xcite showed the vertical relations in word representations across the transformer layers through word identifiability. intuitively, although each word representation has respective knowledge, the set of representations of words as a whole is more semantically meaningful, since words in the embedding space are positioned relatively by learning. inspired by these observations, we propose a novel distillation objective, termed contextual knowledge distillation (ckd), for language tasks that utilizes the statistics of relationships between word representations. in this paper, we define two types of contextual knowledge: word relation (wr) and layer transforming relation (ltr). specifically, wr is proposed to capture the knowledge of relationships between word representations and ltr defines how each word representation changes as it passes through the network layers. we validate our method on general language understanding evaluation (glue) benchmark and the stanford question answer dataset (squad), and show the effectiveness of ckd against the current state-of-the-art distillation methods. to validate elaborately, we conduct experiments on taskagnostic and task-specific distillation settings. we also show that our ckd performs effectively on a variety of network architectures. moreover, with the advantage that ckd has no restrictions on student's architecture, we show ckd further improves the performance of adaptive size pruning method @xcite that involves the architectural changes during the training. to summarize, our contribution is threefold: tang et al. (2019); chia et al. (2019) proposed to distill the bert to train the simple recurrent and convolution networks. sanh et al. (2019); turc et al. (2019) proposed to use the teacher's predictive distribution to train the smaller bert and sun et al. (2019) proposed a method to transfer individual representation of words. in addition to matching the hidden state, jiao et al. (2019); sun et al. (2020); wang et al. (2020) also utilized the attention matrices derived from the transformer. several works including liu et al. (2020); hou et al. (2020) improved the performance of other compression methods by integrating with knowledge distillation objectives in the training procedure. in particular, dynabert (hou et al., 2020) proposed the method to train the adaptive size bert using the hidden state matching distillation. different from previous knowledge distillation methods that transfer respective knowledge of word representations, we design the objective to distill the contextual knowledge contained among word representations. ## setup and background most of the recent state-of-the-art language models are stacking transformer layers which consist of repeated multi-head attentions and position-wise feed-forward networks. transformer based networks. given an input sentence with n tokens, @xmath0, x 2 , . . . , x n ] ∈ r d i ×n , most networks @xcite @xcite utilize the embedding layer to map an input sequence of symbol representations x to a sequence of continuous representations @xmath1, . . . , e n ] ∈ r de×n . then, each l-th transformer layer of the identical structure takes the previous representations r l-1 and produces the updated representations r @xmath2,1 , r l,2 , . . . , r l,n ] ∈ r dr×n through two sublayers: multi-head attention (mha) and positionwise feed forward network (ffn). the input at the first layer (@xmath3. in mha operation where h separate attention heads are operating independently, each input token r l-1,i for each head is projected into a query q @xmath4, key k @xmath5, and value v @xmath6, typically the outputs of all heads are then concatenated and fed through the ffn, producing the single word representation r l,i . for clarity, we pack attention values of all words into a matrix form a l,@xmath7, a 2 , .., a n ] ∈ r n×n for attention head h. knowledge distillation for transformer. in the general framework of knowledge distillation, teacher network (t ) with large capacity is trained in advance, and then student network (s) with pre-defined architecture but relatively smaller than teacher network is trained with the help of teacher's knowledge. specifically, given the teacher parameterized by θ t , training the student parameterized by θ s aims to minimize two objectives: i) the crossentropy loss l ce between the output of the student network s and the true label y and ii) the difference of some statistics l d between teacher and student models. overall, our goal is to minimize the following objective function: ## contextual knowledge distillation we now present our distillation objective that transfers the structural or contextual knowledge which is defined based on the distribution of word representations. unlike previous methods distilling each word separately, our method transfers the information contained in relationships between words or between layers, and provides a more flexible way of constructing embedding space than directly matching representations. the overall structure of our method is illustrated in figure 1(a) . specifically, we design two key concepts of contextual knowledge from language models: word relation-based and layer transforming relation-based contextual knowledge, as shown in figure 1(b) . ## experiments we conduct task-agnostic and task-specific distillation experiments to elaborately compare our ckd with baseline distillation objectives. we then report on the performance gains achieved by our method for bert architectures of various sizes and inserting our objective for training dynabert which can run at adaptive width and depth through pruning the attention heads or layers. finally, we analyze the effect of each component in our ckd and the impact of leveraging locality δ for w ijk in eq. ( 1 ). ## conclusion we proposed a novel distillation strategy that leverages contextual information efficiently based on word relation and layer transforming relation. to our knowledge, we are the first to apply this contextual knowledge which is studied to interpret the language models. through various experiments, we show not only that ckd outperforms the state-ofthe-art distillation methods but also the possibility that our method boosts the performance of other compression methods.