Contrastive Distillation on Intermediate Representations for Language Model Compression (2020)

Abstract: Existing language model compression methods mostly use a simple L_2 loss to distill knowledge in the intermediate representations of a large BERT model to a smaller one. Although widely used, this objective by design assumes that all the dimensions of hidden representations are independent, failing to capture important structural knowledge in the intermediate layers of the teacher network. To achieve better distillation efficacy, we propose Contrastive Distillation on Intermediate Representations (CoDIR), a principled knowledge distillation framework where the student is trained to distill knowledge through intermediate layers of the teacher via a contrastive objective. By learning to distinguish positive sample from a large set of negative samples, CoDIR facilitates the student’s exploitation of rich information in teacher’s hidden layers. CoDIR can be readily applied to compress large-scale language models in both pre-training and finetuning stages, and achieves superb performance on the GLUE benchmark, outperforming state-of-the-art compression methods.

## introduction large-scale pre-trained language models (lms), such as bert @xcite , xlnet @xcite and roberta @xcite , have brought revolutionary advancement to the nlp field @xcite . however, as newgeneration lms grow more and more into behemoth size, it becomes increasingly challenging to deploy them in resource-deprived environment. naturally, there has been a surge of research interest in developing model compression methods @xcite @xcite to reduce network size in pre-trained lms, while retaining comparable performance and efficiency. pkd @xcite was the first known effort in this expedition, an elegant and effective method that uses knowledge distillation (kd) for bert model compression at finetuning stage. later on, distilbert @xcite , tinybert @xcite and mobilebert @xcite carried on the torch and extended similar compression techniques to pre-training stage, allowing efficient training of task-agnostic compressed models. in addition to the conventional kl-divergence loss applied to the probabilistic output of the teacher and student networks, an l 2 loss measuring the difference between normalized hidden layers has proven to be highly effective in these methods. however, l 2 norm follows the assumption that all dimensions of the target representation are independent, which overlooks important structural information in the many hidden layers of bert teacher. motivated by this, we propose contrastive distillation for intermediate representations (codir), which uses a contrastive objective to capture higher-order output dependencies between intermediate representations of bert teacher and the student. contrastive learning @xcite aims to learn representations by enforcing similar elements to be equal and dissimilar elements further apart. formulated in either supervised or unsupervised way, it has been successfully applied to diverse applications @xcite @xcite . to the best of our knowledge, utilizing contrastive learning to compress large transformer models is still an unexplored territory, which is the main focus of this paper. a teacher network's hidden layers usually contain rich semantic and syntactic knowledge that can be instrumental if successfully passed on to the student @xcite @xcite . thus, instead of directly applying contrastive loss to the final output layer of the teacher, we apply contrastive learning to its intermediate layers, in addition to the use of kl-loss between the probabilistic outputs of the teacher and student. this casts a stronger regularization effect for student training by capturing more informative signals from intermediate representations. to maximize the exploitation of intermediate layers of the teacher, we also propose the use of mean-pooled representation as the distillation target, which is empirically more effective than commonly used special [cls] token. to realize constrastive distillation, we define a congruent pair (h t i , h s i ) as the pair of representations of the same data input from the teacher and student networks, as illustrated in figure 1 . incongruent pair (h t i , h s j ) is defined as the pair of representations of two different data samples through the teacher and the student networks, respectively. the goal is to train the student network to distinguish the congruent pair from a large set of incongruent pairs, by minimizing the constrastive objective. for efficient training, all data samples are stored in a memory bank @xcite . during finetuning, incongruent pairs can be selected by choosing sample pairs with different labels to maximize the distance. for pre-training, however, it is not straightforward to construct incongruent pairs this way as labels are unavailable. thus, we randomly sample data points from the same mini-batch pair to form incongruent pairs, and construct a proxy congruent-incongruent sample pool to assimilate what is observed in the downstream tasks during finetuning stage. this and other designs in codir make constrative learning possible for lm compression, and have demonstrated strong performance and high efficiency in experiments. our contributions are summarized as follows. ## related work language model compression to reduce computational cost of training large-scale language models, many model compression techniques have been developed, such as quantization @xcite , pruning @xcite @xcite , knowledge distillation @xcite @xcite @xcite , and direct transformer block modification @xcite . quantization refers to storing model parameters from 32-or 16-bit floating number to 8-bit or even lower. directly truncating the parameter values will cause significant accuracy loss, hence quantizationaware training has been developed to maintain similar accuracy to the original model @xcite . @xcite found that even after most attention heads are removed, the model still retains similar accuracy, indicating there is high redundancy in the learned model weights. later studies proposed different pruning-based methods. for example, gordon et al. (2020) simply removed the model weights that are close to zero; while @xcite used re-weighted l 1 and proximal algorithm to prune weights to zero. note that simple pruning does not improve inference speed, unless there is structure change such as removing the whole attention head. there are also some efforts that try to improve the transformer block directly. typically, language models such as bert @xcite and roberta @xcite can only handle a sequence of tokens in length up to 512. @xcite proposed to use reversible residual layers and locality-sensitive hashing to reduce memory usage to deal with extremely long sequences. besides, @xcite proposed to use convolutional neural networks to capture short-range attention such that reducing the size of self-attention will not significantly hurt performance. another line of research on model compression is based on knowledge transfer, or knowledge distillation (kd) @xcite , which is the main focus of this paper. note that previously introduced model compression techniques are orthogonal to kd, and can be bundled for further speedup. distilled bilstm @xcite tried to dis- trm trm trm trm x x trm trm teacher model (𝑓 𝑡 ) student model (𝑓 𝑠 ) ℎ 0 𝑠 ℎ 0 𝑡 𝑋 0 = the greatest musicians [true] 𝑋 1 = cold movie [false] 𝑋 2 = no apparent joy [false] .... 𝑋 𝐾 = the action is stilted [false] .... ℒ 𝐾𝐷 ℒ 𝐶𝐸 or ℒ 𝑀𝐿𝑀 ℒ 𝐶𝑅𝐷 ℎ 2 𝑠 ℎ 1 𝑠 ℎ 𝐾 𝑠 training data contrastive distillation intermediate representations 𝑋 0 = [mask] music is a [mask] of ... till knowledge from bert into a simple lstm. though achieving more than 400 times speedup compared to bert-large, it suffers from significant performance loss due to the shallow network architecture. distilbert @xcite proposed to distill predicted logits from the teacher model into a student model with 6 transformer blocks. bert-pkd @xcite proposed to not only distill the logits, but also the representation of @xcite tokens from the intermediate layers of the teacher model. tinybert @xcite , mobilebert @xcite and sid @xcite further proposed to improve bert-pkd by distilling more internal representations to the student, such as embedding layers and attention weights. these existing methods can be generally divided into two categories: (i) task-specific, and (ii) task-agnostic. task-specific methods, such as distilled bilstm, bert-pkd and sid, require the training of individual teacher model for each downstream task; while task-agnostic methods such as distilbert, tinybert and mobilebert use kd to pre-train a model that can be applied to all downstream tasks by standard finetuning. contrastive representation learning contrastive learning @xcite is a popular research area that has been successfully applied to density estimation and representation learning, especially in self-supervised setting @xcite . it has been shown that the contrastive objective can be interpreted as maximizing the lower bound of mutual information between different views of the data @xcite @xcite . however, it is unclear whether the success is determined by mutual information or by the specific form of the contrastive loss @xcite . recently, it has been extended to knowledge distillation and cross-modal transfer for image classification tasks @xcite . different from prior work, we propose the use of contrastive objective for transformer-based model compression and focus on language understanding tasks. ## codir for model compression in this section, we first provide an overview of the proposed method in sec. 3.1, then describe the details of contrastive distillation in sec. 3.2. its adaptation to pre-training and finetuning is further discussed in sec. 3.3. ## experiments in this section, we present comprehensive experiments on a wide range of downstream tasks and provide detailed ablation studies, to demonstrate the effectiveness of the proposed approach to largescale lm compression. ## conclusion in this paper, we present codir, a novel approach to large-scale language model compression via the use of contrastive loss. codir utilizes information from both teacher's output layer and its intermediate layers for student model training. extensive experiments demonstrate that codir is highly effective in both finetuning and pre-training stages, and achieves state-of-the-art performance on glue benchmark compared to existing models with a similar size. all existing work either use bertbase or roberta-base as teacher. for future work, we plan to investigate the use of a more powerful language model, such as megatron-lm @xcite , as the teacher; and different strategies for choosing hard negatives to further boost the performance.