Distiller: A Systematic Study of Model Distillation Methods in Natural Language Processing (2021)

Abstract: Knowledge Distillation (KD) offers a natural way to reduce the latency and memory/energy usage of massive pretrained models that have come to dominate Natural Language Processing (NLP) in recent years. While numerous sophisticated variants of KD algorithms have been proposed for NLP applications, the key factors underpinning the optimal distillation performance are often confounded and remain unclear. We aim to identify how different components in the KD pipeline affect the resulting performance and how much the optimal KD pipeline varies across different datasets/tasks, such as the data augmentation policy, the loss function, and the intermediate representation for transferring the knowledge between teacher and student. To tease apart their effects, we propose Distiller, a meta KD framework that systematically combines a broad range of techniques across different stages of the KD pipeline, which enables us to quantify each component’s contribution. Within Distiller, we unify commonly used objectives for distillation of intermediate representations under a universal mutual information (MI) objective and propose a class of MI-objective functions with better bias/variance trade-off for estimating the MI between the teacher and the student. On a diverse set of NLP datasets, the best Distiller configurations are identified via large-scale hyper-parameter optimization. Our experiments reveal the following: 1) the approach used to distill the intermediate representations is the most important factor in KD performance, 2) among different objectives for intermediate distillation, MI-performs the best, and 3) data augmentation provides a large boost for small training datasets or small student networks. Moreover, we find that different datasets/tasks prefer different KD algorithms, and thus propose a simple AutoDistiller algorithm that can recommend a good KD pipeline for a new dataset.

## introduction recent advancements in natural language processing (nlp) such as bert @xcite , roberta @xcite , and electra @xcite have demonstrated the effectiveness of transformer models in generating transferable language representations. pretraining over a massive unlabeled corpus and then fine-tuning over labeled data for the task of interest has become the state-ofthe-art paradigm for solving diverse nlp problems ranging from sentence classification to question answering @xcite . scaling up the size of these networks has led to rapid nlp improvements. however, these improvements have come at the expense of significant increase of memory for their many parameters and compute to produce predictions @xcite . this prevents these models from being deployed on resource-constrained devices such as smart phones and browsers, or for latency-constrained applications such as click-through-rate prediction. there is great demand for models that are smaller in size, yet still retain similar accuracy as those having a large number of parameters. to reduce the model size while preserving accuracy, various model compression techniques have been proposed such as: pruning, quantization, and knowledge distillation (kd) @xcite . among these methods, task-aware kd is a popular and particularly promising approach for compressing transformer-based models @xcite . the general idea is to first fine-tune a large model (namely the teacher model) based on the task's labeled data, and then train a separate network that has significantly fewer parameters (namely the student model) than the original model to mimic the predictions of the original large model. a large number of task-aware kd algorithms have been proposed in the nlp regime, e.g., distillbert @xcite , bert-pkd @xcite , bert-emd @xcite , tiny-bert @xcite , dynabert @xcite , and adabert @xcite , some of which can compress the teacher network by 10× without significant accuracy loss on certain datasets. innovations in kd for nlp generally involve improvements in one of the following aspects: 1) the loss function for gauging the discrepancy between student and teacher predictions @xcite , 2) the method for transferring intermediate network representations between teacher and student @xcite @xcite , 3) the use of data augmentation during student training @xcite , and 4) multiple stages of distillation @xcite . many research proposals have simultaneously introduced new variations of more than one of these components, which confounds the impact of each component on the final performance of the distillation algorithm. in addition, it is often unclear whether a proposed distillation pipeline will generalize to a new dataset or task, making automated kd challenging. for example, mixkd @xcite has only been evaluated on classification problems and it is unclear if the method will be effective for question answering. to understand the importance of different components in kd, we undertake a systematic study of kd algorithms in nlp. our study is conducted using a meta-distillation pipeline we call distiller that contains multiple configurable components. all candidate algorithms in the search space of distiller work for two types of nlp tasks: text classification and sentence tagging. distiller unifies existing techniques for knowledge transfer from intermediate layers of the teacher network to the student network (i.e. intermediate distillation) as special cases of maximizing (bounds on) the mutual information (mi) between teacher and student representations. based on this unification and recent progress in variational bounds of mi @xcite , we propose a new intermediate distillation objective called mi-α that uses a scalar α to control the bias-variance trade-off of mi estimation. including mi-α in the search space of distiller, we run extensive hyper-parameter tuning algorithms to search for the best distiller-configuration choices over glue @xcite and squad @xcite . this search helps us identify the best distillation pipelines and understand what impact different kd modules have on student performance in nlp. using the observations of this large-scale study, we train a autodistiller model to predict the distillation ratio, which is defined as the fraction of the teacher's performance achieved by the student, based on kd pipeline choices and characteristics of a dataset. leave-one-out cross validation evaluation of autodistiller demonstrates that it is able to reliably prioritize high-performing kd configurations in most folds, and is able to suggest good distillation pipelines on two new datasets. the main contributions of this work include: knowledge distillation. the general kd framework was popularized by buciluǎ et al. (2006); hinton et al. (2014), aiming to transfer knowledge from an accurate but cumbersome teacher model to a compact student model by matching the class probabilities produced by the teacher and the student. focusing on automl settings with tabular data, fakoor et al. (2020) proposed a general kd algorithm for different classical ml models and ensembles thereof. also hoping to identify good choices in the kd pipeline like our work, kim et al. (2021) compared kullback-leibler divergence and mean squared error objectives in kd for image classification models, finding that mean squared error performs better. recent kd research in the domain of nlp has investigated how to efficiently transfer knowledge from pretrained transformer models. sun et al. (2019) proposed bert-pkd that transfers the knowledge from both the final layer and intermediate layers of the teacher network. jiao et al. (2020) proposed the tinybert model that first distills the general knowledge of the teacher by minimizing the masked language model (mlm) objective (devlin et al., 2019), with subsequent task-specific distillation. li et al. (2020) proposed a many-to-many layer mapping function leveraging the earth mover's distance to transfer intermediate knowledge. our paper differs from the existing work in that we provide a systematic analysis of the different components in single-stage task-aware kd algorithms in nlp and propose the first automated kd algorithm in this area. ## methodology our study is structured around a configurable metadistillation pipeline called distiller. distiller contains four configurable components, namely: a data augmentation policy a(•, •), a layer mapping configuration of intermediate distillation {m i,j }, an intermediate distillation objective l inter (•, •), and a prediction layer distillation objective l pred (•, •). assume the teacher network f t has m layers and the student network f s has n layers. then for a given data/label pair (x, y) sampled from the dataset d, the student acquires knowledge from the teacher by minimizing the following objective: here m i,@xmath0, 1] represents the layer mapping weight between the i-th teacher layer and j-th student layer, h t i , h s j are the i-th and the j-th hidden states of the teacher and the student (i.e. their intermediate representations at layers i and j), β 1 , β 2 control the strength of distilling from class probabilities produced by the teacher, and γ 1 and γ 2 control the strength of learning from ground truth data (x, y) and synthesized (augmented) data (x, ŷ). in appendix, we illustrated how previous model distillation algorithms @xcite @xcite can be encompassed in the distiller framework. ## experimental setup to study the importance of each component described in the previous section, we randomly sample distiller configurations in the designed search space while fixing the optimizer and other unrelated hyper-parameters. we apply each sampled distillation configuration on a diverse set of nlp tasks and different teacher/student architectures. all experiments are evaluated on glue @xcite and squad v1.1 @xcite that contain classification, regression, and sentence tagging tasks. here, we view the question answering problem in squad v1.1 as finding the correct answer span from the given context, which is essentially a sentence tagging task. we adopt the same metrics for these tasks as in the original papers @xcite . since @xcite finds initializing students with pretrained weights is better for distillation, we initialize student models with either weights obtained from taskagnostic distillation @xcite or pretrained from scratch @xcite . three different pretrained models bert base @xcite , roberta large @xcite and electra large @xcite are considered as teacher models in our experiments after task-specific fine-tuning. as student models, we consider options like tinybert 4 , electra small , as well as other models detailed in existing implementations of data augmentation in kd for nlp generally first generate an augmented dataset that is k times larger than the original one and then apply the distillation algorithm over the augmented dataset. such implementation separates the process of da from kd, leading to a time/storage-consuming and inflexible kd pipeline. in distiller, we instead apply da dynamically during training. in addition, we use the teacher network to compute the soft label ŷ assigned to any augmented sample x. to analyze the importance of different components in distiller, we adopted fanova @xcite , an algorithm for quantifying the importance of individual hyper-parameters as well as their interactions in determining down-stream performance. we use fanova to evaluate the importance of the four components in distiller as well as their pairwise combinations: data augmentation, intermediate distillation objective, layer mapping strategy, and prediction layer distillation objective. ## experimental results under the previously described experimental setup, we conducted a random search over distiller configurations on each dataset and collected more than 1300 data points in total. each collected data point contains a particular distiller configuration, the dataset/task, the teacher/student architectures, and the final performance of the student model. analyzing the data reveals three major findings: ## conclusion we provided a systematic study of kd algorithms in nlp to understand the importance of different components in the kd pipeline for various nlp tasks. using data collected from our study, we fit a autodistiller to predict student performance under each kd pipeline based on dataset features, which is helpful for automatically selecting which kd pipeline to use for a new dataset. here we unified the existing intermediate distillation objectives as maximizing the lower bounds of mi, leading to a new mi-α objective based on tighter bounds, which performs better on many datasets.