How to Distill your BERT : An Empirical Study on the Impact of Weight Initialisation and Distillation Objectives (2023)

Abstract: Recently, various intermediate layer distillation (ILD) objectives have been shown to improve compression of BERT models via Knowledge Distillation (KD). However, a comprehensive evaluation of the objectives in both task-specific and task-agnostic settings is lacking. To the best of our knowledge, this is the first work comprehensively evaluating distillation objectives in both settings. We show that attention transfer gives the best performance overall. We also study the impact of layer choice when initializing the student from the teacher layers, finding a significant impact on the performance in task-specific distillation. For vanilla KD and hidden states transfer, initialisation with lower layers of the teacher gives a considerable improvement over higher layers, especially on the task of QNLI (up to an absolute percentage change of 17.8 in accuracy). Attention transfer behaves consistently under different initialisation settings. We release our code as an efficient transformer-based model distillation framework for further studies.

## introduction large-scale pre-trained language models (plms) have brought revolutionary advancements to natural language processing, such as bert @xcite , xlnet @xcite , electra @xcite and gpt-3 @xcite . however, the enormous size of these models has led to difficulties in deploying them in resourceconstrained environments. therefore significant interest has emerged in developing methods for reducing their size. knowledge distillation (kd) @xcite transfers the knowledge embedded in one model to another, which can be used for crosslingual transfer, cross-modal transfer, and model compression. kd heavily depends on the distillation objective, which determines how knowledge 1 https://github.com/mainlp/how-to-distill-your-bert is transferred. many works have tried to design different distillation objectives for transformerbased @xcite model compression and successfully distilled plms into smaller models, either task-specifically @xcite or task-agnostically-which differ in whether kd is performed at the pre-training stage or during task finetuning @xcite @xcite . despite their impressive results, determining the best distillation objective is difficult due to their diverse comparison setups, such as data preprocessing, student model initialization, layer mapping strategies, task-specific/agnostic settings, and others. this breadth of choices and lack of code has led to comparison on unequal grounds and contradictory findings. 2 this shows a substantial need to reproduce and evaluate distillation objectives within the same setting. motivated by this gap, we conduct experiments on the most common distillation objectives and their combinations in taskspecific and task-agnostic settings. from our empirical evaluation, we show: (1) attention transfer performs consistently well in various initialisation settings, (2) initialisation with lower layers of the teacher gives a considerable improvement over higher layers in task-specific distillation. in summary, our contributions are: ## distillation objectives prediction layer transfer prediction layer transfer minimizes the soft cross-entropy between the logits from the teacher and the student: l @xmath0, z s /t , with z t and z s the logits from the teacher/student and t is the temperature value. following the vanilla kd approach @xcite , the final training loss is a combination of l pred and supervision loss l ce (masked language modelling loss l mlm in the pertaining stage). we denote this objective as vanilla kd. hidden states transfer hidden states transfer penalizes the distance between the hidden states of specific layers from the teacher and the student. common choices for the representation are the embedding of the [cls] token @xcite and the whole sequence embedding @xcite . we use mean-squared-error (mse) to measure the distance between the student and teacher embedding, which can be formulated as l @xmath1, h t , where h h h @xmath2, d and d ′ are the hidden dimensions. the matrix w w w @xmath3. we denote this objective as hid-cls. in the case of transferring the sequence embedding, one can replace the token embeddings with sequence embeddings h h h @xmath4, where l is the sequence length. the objective that transfers the sequence embedding with mse loss is denoted as hid-seq. we also evaluated a contrastive representation learning method which transfers the hidden state representation from the teacher to the student with a contrastive objective @xcite . we inherited their code for implementation and refer our readers to the original paper for details. we denote this objective as hid-cls-contrast. attention and value transfer the attention mechanism has been found to capture rich linguistic knowledge @xcite , and attention map transfer is widely used in transformer model distillation. to measure the similarity between the multi-head attention block of the teacher and the student, mse and kullback-leibler divergence are the two standard loss functions. the objective using mse is formulated as since the attention after the softmax function is a distribution over the sequence, we can also use the kl-divergence to measure the distance: ## experimental setup we evaluate our model on the general language understanding evaluation (glue) benchmark @xcite tasks, including linguistic acceptability (cola), sentiment analysis (sst-2), semantic equivalence (mrpc, qqp), and natural language inference (mnli, qnli, rte). for task-specific distillation, we distill a finetuned roberta base @xcite into a 3layer transformer model on each glue task, using the fairseq @xcite implementation and the recommended hyperparameters presented in @xcite . we follow the training procedure from tinybert to perform intermediate layer and prediction layer distillation sequentially for 10 epochs each, freeing us from tuning the loss weights. for intermediate layer distillation, the student learns from the same teacher's layers that were used for initialising the student. in addition, we always initialise the embedding layer with the teacher's embedding layer. for task-agnostic distillation, we distill the uncased version of bert base into a 6-layer student model, based on the implementation by @xcite . here we perform last-layer knowledge transfer since we see no improvement when transferring multiple layers in our experiments. we train the student model for 100k steps with batch size 1024, a peaking learning rate of 5e-4 and a maximum sequence length of 128. the distilled student model is then fine-tuned on the glue datasets with grid search over batch size {16, 32} and learning rate {1e-5, 3e-5, 5e-5, 8e-5}. we follow the original training corpus of bert: english wikipedia and bookcorpus @xcite . for vanilla kd and hid-cls transfer, the performance on qnli is significantly improved when initialising with lower teacher layers. attention transfer benefits less from initialising from lower teacher layers. ## results distillation objectives distillation objective performances are compared in in the task-agnostic setting, the att-mse objective outperforms att-kl, which performs similarly to vanilla kd and hidden states transfer. this contradicts the observation in minilm @xcite , where their att-kl based objective outperforms tinybert @xcite with att-mse. however, minilm has more training iterations and a larger batch size, which makes comparison difficult. the performance drop of att-kl compared to att-mse is mainly due to its poor performance on cola (linguistic acceptability of a sentence), on which minilm also performs poorly. we hypothesise that mse can transfer the linguistic knowledge embedded in the attention matrix more effectively because the mse loss function gives more direct matching than kl-divergence, which was also concluded by @xcite . for reference, we report the result of 3 existing works that use the same objectives in our experiments. the result of distilbert and minilm are taken from the respective papers. the result of tinybert is taken from @xcite for fair comparison since tinybert only reported taskspecific distillation result with data augmentation. we denote the prior works and the corresponding objective we evaluate with the same superscript symbol. initialisation we also studied the impact of the choice of teacher layers for initialising the student. evaluation score on glue task development sets under different teacher layer choices for initialisation are reported in in the task-agnostic setting, we only observe considerable improvement with the objective hid-cls, which performs poorly when randomly initialized, compared to other objectives. this contradicts @xcite with a vanilla kd objective, where they instead showed improvement of 3 average score when initialising from the teacher over random initialisation. however, our vanilla-kd approach initialised with random weights outperforms their best result (79.3 vs 78.5). therefore, we hypothesise that the advantage of pre-loading teacher layers over random initialisation diminishes as the student is fully distilled during pre-training. ## conclusion we extensively evaluated distillation objectives for the transformer model and studied the impact of weight initialisation. we found that attention transfer performs consistently well in both task-specific and task-agnostic settings, regardless of the teacher layers chosen for student initialization. we also observed that initialising with lower teacher layers significantly improved task-specific distillation performance compared to higher layers. we release our code and hope this work motivates further research into developing better distillation objectives and compressing in-house models.