CILDA : Contrastive Data Augmentation Using Intermediate Layer Knowledge Distillation (2022)

Abstract: Knowledge distillation (KD) is an efficient framework for compressing large-scale pre-trained language models. Recent years have seen a surge of research aiming to improve KD by leveraging Contrastive Learning, Intermediate Layer Distillation, Data Augmentation, and Adversarial Training. In this work, we propose a learning-based data augmentation technique tailored for knowledge distillation, called CILDA. To the best of our knowledge, this is the first time that intermediate layer representations of the main task are used in improving the quality of augmented samples. More precisely, we introduce an augmentation technique for KD based on intermediate layer matching using contrastive loss to improve masked adversarial data augmentation. CILDA outperforms existing state-of-the-art KD approaches on the GLUE benchmark, as well as in an out-of-domain evaluation.

## introduction the exponentially increasing size of pre-trained large language models @xcite @xcite has been a persistent concern regarding the efficiency and scalability of natural language understanding (nlu) in real world applications. knowledge distillation (kd) @xcite is a technique for transferring the knowledge from a large-scale model (called teacher) to a smaller one (called student), so that the latter model can be employed on edge device @xcite @xcite . this is done by minimizing the kl divergence between the teacher and student probabilistic outputs. numerous techniques have been exploited recently to increase the knowledge transfer beyond logits matching. for instance, it has been found beneficial to perform distillation on the internal components (parameters) of the teacher and student, which is known as intermediate layer distillation @xcite @xcite @xcite . data augmentation has also been successful for kd @xcite @xcite , as researchers have found that the student has less opportunity to acquire useful information from the teacher when limited data are available for training @xcite @xcite . adversarial training was also employed in kd @xcite @xcite to improve the robustness and generalization, as the student may predict inconsistent outputs with slight distortion to the data distributions @xcite . recently, contrastive learning @xcite @xcite has been exploited for improving knowledge transfer @xcite , and to optimize the intermediate layer mapping scheme @xcite . each of the aforementioned techniques has proven effective in addressing a specific challenge in kd. yet, we are not aware of a single method that takes advantage of all of them. in this paper, we propose cilda, a kd method that incorporate contrasting learning, intermediate layer distillation, data augmentation, and adversarial training. distilling into a 6-layer bert model, cilda delivers new state-of-the-art results on the glue benchmark @xcite , as well as outperforming other kd methods in out-of-domain evaluations. ## related work many studies @xcite @xcite have noticed that important structural linguistic information are hidden in the intermediate layers of transformer models @xcite . recent kd methods propose to match teacher and student: intermediate layers representations @xcite @xcite , embedding matrix @xcite , and self-attention distributions @xcite . other variants of kd methods have been proposed such as annealing-kd @xcite and pro-kd @xcite , two stage distillation methods where a smooth and gradual training of the student is controlled by a dynamic temperature factor, followed by a simple cross entropy loss for a few epochs. augmented adversarial examples @xcite are label-preserving transformations in the embedding space that are used to improve generalizability of models. freelb @xcite is an adversarial algorithm which creates virtual adversarial examples from word embeddings, and then performs the parameter updates on these adversarial embeddings. mate-kd @xcite is a min-max adversarial data augmentation approach for kd, where an extra generator model is trained to generate adversarial text by maximizing the logit output margins between the teacher and the student. contrastive learning is a self-supervised representation learning method @xcite @xcite which learns the feature representation of the samples by contrasting positive and negative samples. codir @xcite ) is a contrast-enhanced diversity promoting method between teacher and student intermediate representations of data samples from the same class. mate-kd is the most related to our solution, with one notable difference: we believe our technique is the first to deploy intermediate layers distillation with the contrastive objective in the data augmentation process. ## cilda in this section, we introduce cilda, our contrastive approach for masked adversarial text augmentation for knowledge distillation using intermediate layer matching. inspired by @xcite , we deploy a generator (e.g. bert) which will be trained to map masked inputs, x, to augmented samples, x . the objective of this mapping is to perturb the inputs (in their vicinity) such that their corresponding output and intermediate layer representations of the teacher and student networks diverge to their maximum. generating such maximum divergence augmented samples aims to fill the existing major gaps in the training data. we mask input tokens with a certain pre-defined probability, p. the architecture of our model is depicted in figure 1. our training is comprised of two alternating steps we describe hereafter. ## conclusion and future work we proposed a min-max adversarial data augmentation framework for kd, which is powered by contrastive distillation loss for intermediate layer matching. our algorithm maximizes the intermediate and logit representation margin between the teacher and the student models. in future works, we would like to investigate the distillation from superlarge models such as megatron @xcite and t5 @xcite . also, we would like to improve the generator output quality via distillation from generative models like @xcite .