MATE - KD : Masked Adversarial TE xt, a Companion to Knowledge Distillation (2021)

Abstract: The advent of large pre-trained language models has given rise to rapid progress in the field of Natural Language Processing (NLP). While the performance of these models on standard benchmarks has scaled with size, compression techniques such as knowledge distillation have been key in making them practical. We present MATE-KD, a novel text-based adversarial training algorithm which improves the performance of knowledge distillation. MATE-KD first trains a masked language model-based generator to perturb text by maximizing the divergence between teacher and student logits. Then using knowledge distillation a student is trained on both the original and the perturbed training samples. We evaluate our algorithm, using BERT-based models, on the GLUE benchmark and demonstrate that MATE-KD outperforms competitive adversarial learning and data augmentation baselines. On the GLUE test set our 6 layer RoBERTa based model outperforms BERT-large.

## introduction transformers @xcite and transformer-based pre-trained language models (plms) @xcite are ubiquitous in applications of nlp. they are highly parallelizable and their performance scales well with an increase in model parameters and data. increasing model parameters depends on the availability of computational resources and plms are typically trained on unlabeled data which is cheaper to obtain. recently, the trillion parameter mark has been breached for plms @xcite amid serious environmental concerns @xcite . however, without a change in our current training paradigm , training larger models may be unavoidable @xcite . in order to deploy these models for practical applications such as for virtual personal assistants, recommendation systems, e-commerce platforms etc. model compression is necessary. knowledge distillation (kd) @xcite is a simple, yet powerful knowledge transfer algorithm which is used for neural model compression @xcite , ensembling @xcite and multi-task learning @xcite . in nlp, kd for compression has received renewed interest in the last few years. it is one of the most widely researched algorithms for the compression of transformer-based plms @xcite . one key feature which makes kd attractive is that it only requires access to the teacher's output or logits and not the weights themselves. therefore, if a trillion parameter model resides on the cloud, an api level access to the teacher's output is sufficient for kd. consequently, the algorithm is architecture agnostic, i.e., it can work for any deep learning model and the student can be a different model from the teacher. recent works on kd for transfer learning with plms extend the algorithm in two main directions. the first is towards "model" distillation @xcite @xcite i.e. distilling the intermediate weights such as the attention weights or the intermediate layer output of transformers. the second direction is towards curriculum-based or progressive kd @xcite @xcite where the student learns one layer at a time or from an intermediary teacher, known as a teacher assistant. while these works have shown accuracy gains over standard kd, they have come at the cost of architectural assumptions, least of them a common architecture between student and teacher, and greater access to teacher parameters and intermediate outputs. another issue is that the decision to distill one teacher layer and to skip another is arbitrary. still the teacher typically demonstrates better generalization we are interested in kd for model compression and study the use of adversarial training @xcite to improve student accuracy using just the logits of the teacher as in standard kd. specifically, our work makes the following contributions: ## methodology we propose an algorithm that involves co-training and deploy an adversarial text generator while training a student network using kd. figure 1 gives an illustration of our architecture. ## experiments we evaluated mate-kd on all nine datasets of the general language understanding evaluation (glue) @xcite benchmark which include classification and regression. these datasets can be broadly divided into 3 families of problems. single set tasks which include linguistic acceptability (cola) and sentiment analysis (sst-2). similarity and paraphrasing tasks which include paraphrasing (mrpc and qqp) and a regression task (sts-b). inference tasks which include natural language inference (mnli, wnli, rte) and question answering (qnli). ## discussion and future work we have presented mate-kd, a novel text-based adversarial training algorithm which improves the student model in kd by generating adversarial examples while accessing the logits of the teacher only. this approach is architecture agnostic and can be easily adapted to other applications of kd such as model ensembling and multi-task learning. we demonstrate the need for an adversarial training algorithm for kd based on text rather than embedding perturbation. moreover, we demonstrate the importance of masking for our algorithm. one key theme that we have presented in this work is that as plms inevitably increase in size and number of parameters, techniques that rely on access to the various layers and intermediate parameters of the teacher will be more difficult to train. in contrast, algorithms which are wellmotivated and require minimal access to the teacher may learn from more powerful teachers and would be more useful. an example of such an algorithm is the kd algorithm itself. future work will consider a) using label information and a measure of semantic quality to filter the generated sentences b) exploring the application of our algorithm to continuous data such as speech and images and c) exploring other applications of kd.