Enhancing Task-Specific Distillation in Small Data Regimes through Language Generation (2022)

Abstract: Large-scale pretrained language models have led to significant improvements in Natural Language Processing. Unfortunately, they come at the cost of high computational and storage requirements that complicate their deployment on low-resource devices. This issue can be addressed by distilling knowledge from larger models to smaller ones through pseudo-labels on task-specific datasets. However, this can be difficult for tasks with very limited data. To overcome this challenge, we present a novel approach where knowledge can be distilled from a teacher model to a student model through the generation of synthetic data. For this to be done, we first fine-tune the teacher and student models, as well as a Natural Language Generation (NLG) model, on the target task dataset. We then let both student and teacher work together to condition the NLG model to generate examples that can enhance the performance of the student. We tested our approach on two data generation methods: a) Targeted generation using the Monte Carlo Tree Search (MCTS) algorithm, and b) A Non-Targeted Text Generation (NTTG) method. We evaluate the effectiveness of our approaches against a baseline that uses the BERT model for data augmentation through random word replacement. By testing this approach on the SST-2, MRPC, YELP-2, DBpedia, and TREC-6 datasets, we consistently witnessed considerable improvements over the word-replacement baseline.

## introduction transformer-based models have shown wide success over a variety of natural language processing (nlp) tasks. their ability to scale up to trillions of * corresponding author parameters made it possible to acquire and transfer generalized knowledge from large collections of data to downstream tasks. while these models can lead to significant improvements in performance, the increasing size of their learning parameters results in greater complexity and storage requirements. this can be challenging in real-time applications, especially on devices with limited computational resources @xcite . hence, reducing the size of these language models without sacrificing too much of their performance has become the focus of many works in knowledge distillation (kd). instead of optimizing compressed models for hard-labeled data, @xcite proposed to train a smaller model (the student) with the task of predicting the probability distribution outputs (soft labels) from a larger model (the teacher). this approach has been shown to produce comparable results between the student and teacher models, but usually relies on a large enough dataset through which knowledge can be transferred. to help improve the student's learning in kd, large unlabeled datasets are required @xcite . although unlabeled data is cheaper to obtain and is widespread when compared to labeled data, it may not be available for every task and application. we therefore propose to generate synthetic examples that can be used to transfer knowledge to the student in downstream tasks. to overcome the challenges that come with unavailability of large unlabeled datasets required for the distillation process, we build a data generation framework where the monte carlo tree search (mcts) algorithm is applied to help generate examples that, if added to the student's training data, will increase its performance. by taking the difference between the student's and the teacher's uncertainty for a generated example, we are able to force the language generation model to create examples that can be pseudo-labeled by the teacher with higher confidence than its student. we make the assumption that the wider the gap between the student's and teacher's uncertainty for a particular example, the more likely it is that this example is correctly pseudo-labeled by the teacher and the less likely that it is known to the student model. by training the student on the generated data with the teacher's pseudo labels, it is able to improve its performance by mimicking its teacher's behavior. to generate examples that meet this condition, we take advantage of mcts's tendency to search for paths that maximize the reward value, hence, the uncertainty gap. the intuition here is that the larger the positive difference in the uncertainty, the greater the contradiction is between the student and its teacher, and the more likely that the generated example is important for the student's learning stage. we also show that strong results can still be obtained with a random generation approach that does not optimize for a reward function during the generation process. instead, it first generates data, then selects samples that meet the conditions that are set in the reward function. the contributions of the paper are: the remainder of the paper is structured as follows: section 2 provides a background and an overview of related literature. section 3 describes the proposed approach. section 4 presents the experiments which were carried out. section 5 gives conclusions and plans for future work. ## background in the pursuit of improving performance for natural language processing, pretraining large-scale models on increasing amounts of unlabeled data has become a common approach @xcite @xcite . by leveraging the knowledge gained from the pretraining step, these models have shown impressive performance on many downstream text tasks, e.g. glue and superglue benchmarks @xcite . however, such improvements are accompanied by an increasing number of learning parameters, creating a need for more computational and storage requirements. to alleviate the aforementioned complexity issues, many have suggested approaches for efficient training through model optimization e.g. removal of inefficient or redundant parameters @xcite , and knowledge distillation @xcite @xcite . in knowledge distillation, the unlabeled data plays an intermediary role, which allows the teacher to transfer its knowledge through its predictions. when this data lacks the components for a meaningful knowledge transfer, e.g. limitations in size or textual variations, augmentation techniques can be applied to create additional examples. for instance, @xcite ; @xcite apply task-agnostic heuristics like word replacements, to improve distillation on downstream tasks. the concept of augmenting training examples has been successfully shown to improve training in computer vision @xcite , and has been gaining traction in the nlp domain as well. this includes word manipulations such as the deletion, insertion, or addition of random words in text @xcite , paraphrasing or back-translation @xcite , and most recently the application of language models to predict alternative words @xcite . in this work, instead of relying on the above augmentation techniques, we propose to improve knowledge transfer by involving the student and the teacher in the creation of useful examples. we achieve this through a framework that uses an nlg model to create examples that are deemed useful for knowledge distillation. the steps taken to determine the usefulness of an example are summarized in figure 1 . in mctg, we achieve this by exploring mcts's ability for finding optimal solutions which are rewarded by the usefulness of the examples they represent, as explained in section 3.2. ## approach in this work, we propose the monte carlo text generation (mctg) method (section 3.2), alongside the non-targeted text generation (nttg) method (section 3.4). in mctg, we combine mcts, a language generation model, a teacher, and a student classifier to create synthetic examples that can enhance the performance of the student in a kd generated examples gpt-2 data generation components reward function reference model target model tk 3 bos tk 3 tk 1 eos tk 1 mcts tk 1 tk 2 tk 2 figure 2: mcts builds a tree from token sequences generated by gpt-2. meanwhile, the teacher and student models work together to reward for completed paths. the main components of our framework, as shown in figure 2 , include the language generation model (gpt-2), a teacher model, a student model, and the mcts algorithm. below we discuss the role of each component. ## conclusion in this paper, we presented an approach for generating text data in order to improve knowledge distillation on small datasets. by selecting examples predicted with the lowest uncertainty by the teacher and the highest uncertainty by the student, we were able to improve the student's performance, sometimes almost to the level of the teacher. considering the results, we could argue that reward-based language generation can complement or even substitute for heuristic data augmentation approaches in knowledge distillation. we believe that our approach can serve as a baseline for reward-based textual data generation in small data settings. this will hopefully motivate future research to further explore reward-based generation methods.