A Systematic Study of Knowledge Distillation for Natural Language Generation with Pseudo-Target Training (2023)

Abstract: Modern Natural Language Generation (NLG) models come with massive computational and storage requirements. In this work, we study the potential of compressing them, which is crucial for real-world applications serving millions of users. We focus on Knowledge Distillation (KD) techniques, in which a small student model learns to imitate a large teacher model, allowing to transfer knowledge from the teacher to the student. In contrast to much of the previous work, our goal is to optimize the model for a specific NLG task and a specific dataset. Typically in real-world applications, in addition to labeled data there is abundant unlabeled task-specific data, which is crucial for attaining high compression rates via KD. In this work, we conduct a systematic study of task-specific KD techniques for various NLG tasks under realistic assumptions. We discuss the special characteristics of NLG distillation and particularly the exposure bias problem. Following, we derive a family of Pseudo-Target (PT) augmentation methods, substantially extending prior work on sequence-level KD. We propose the Joint-Teaching method, which applies word-level KD to multiple PTs generated by both the teacher and the student. Finally, we validate our findings in an extreme setup with no labeled examples using GPT-4 as the teacher. Our study provides practical model design observations and demonstrates the effectiveness of PT training for task-specific KD in NLG.

## introduction modern natural language generation (nlg) systems are based on pre-trained language models (lms), which are gradually achieving remarkable milestones @xcite @xcite . alongside the impressive advances in applications such as neural machine transla-tion (nmt), summarization, chatbots, such models have also become increasingly larger, deeper, slower, and more complex. the massive storage requirements and high computational complexity of nlg models discourage their deployment in real-life. as such, there is a growing demand in the industry for compressing such models while preserving their performance. model compression methods typically either prune less informative parameters @xcite or use knowledge distillation (kd) @xcite to transfer knowledge from a larger model (the teacher) to a smaller model (the student). in generation tasks, kd can be applied at the word-level, by training the student to mimic the teacher's next token distribution, or at the sequence-level, by training the student on pseudo-targets (pts) generated by the teacher. although kd research is extensive @xcite @xcite , most works focus on natural language understanding (nlu) tasks, task-agnostic language modeling, or specific generation tasks (e.g., nmt). additionally, kd works for nlg typically consider large datasets with hundreds of thousands of labeled examples, and ignore unlabeled data @xcite @xcite . in more realistic scenarios, however, the number of labeled examples is limited, alongside an abundance of unlabeled data @xcite that may contribute to kd. to bridge these gaps, in this paper we conduct a systematic study of kd for nlg, considering a variety of tasks: summarization, question generation, abductive reasoning, style transfer and simplification, in a more realistic setup. our realistic setup follows 5 criteria that are particularly attractive for a broad range of nlp practitioners: (1) only several thousand labeled examples are available for training (medium-resource), as annotation is costly or labor-intensive, especially for nlg. this is in contrast to research setups where labeled datasets can be very large. (2) large amounts of unlabeled data are available, as is often the case in industrial setups where unlabeled data is collected during the life-cycle of the product; (3) off-the-shelf models are used, which is more practical than training models from scratch; (4) inference-time efficiency is our goal, meaning high compression rate; (5) one-time computational training resources are negligible, compared to inference-time, allowing extensive use of pts. recently, huge lms with excellent generative capacity, such as gpt-4 (openai, 2023) have been presented. while it is tempting to focus our research on them, we focus on small to medium size lms in a fine-tuning setup. this choice is because utilizing a huge lm as a teacher is often infeasible, e.g., due to their high financial costs or when the data cannot be sent to external servers because of privacy constraints. furthermore, research suggests that using mediator-teachers aids the distillation process @xcite , as might be the case in distillation from a huge lm to a medium fine-tuned teacher and finally to a small student. for an extended discussion, see §7. our work hence focuses on a medium size finetuned teacher and we assume there are several thousand labeled examples for its fine-tuning. despite the above limitations, applying huge lms in some valuable setups is still possible. therefore, we also consider the distillation of one such model (gpt-4), although this is not our main focus. we start our study by comparing architectural (encoder-decoder vs. decoder-only), pruning and kd design decisions, discussing the tradeoff between computational resources and task performance. we focus on practical measures like latency and throughput, which is important for batchoffline applications and is typically overlooked. we next provide the first exposure bias perspective for kd which motivates pt augmentation. this bias derives from teacher forcing when the lm conditions on ground-truth tokens during training, while at inference time it conditions on previously generated tokens @xcite . as the distillation progresses, the student's predictions gradually become similar to its teacher's, and therefore training with pts can alleviate exposure bias. we propose extensions of the common practice of generating a single mode approximation pt via beam search, instead, we suggest sampling multiple pts to facilitate higher exposure to conditional distribution factors. additionally, we generate pts for unlabeled data and demonstrate their effectiveness. moreover, we propose a novel kd technique termed joint-teaching, which applies word-level kd to pts generated by both the teacher and the student. this technique aims to implicitly and explicitly address the student exposure bias, ground the learning and teach it to correct its mistakes. finally, we extend the scope of our study by working solely with limited unlabeled examples. due to the absence of labeled examples, fine-tuning the teacher is infeasible, leading us to depend on huge lms with zero-shot capabilities. consequently, we investigate whether our kd findings from our realistic setup (which involves a finetuned teacher) remain applicable to the new extreme setup. to this end, we show how to successfully distill gpt-4, a huge decoder-only model, into a small encoder-decoder model (t5-small), which also has a different tokenizer. our main empirical findings ( §5) are: (1) encoder-decoder architectures outperform their decoder-only counterparts in task-specific finetuning for nlg; (2) decoder pruning substantially outperforms encoder pruning when considering both latency and task performance; and (3) pts can be used much more effectively compared to what was suggested in previous work and this yields substantially improved task performance on a much reduced computational cost. ## experimental setup in this section we describe our four nlg tasks and datasets, the participating models and the evaluation procedures. urls of the code and datasets, as well as implementation details and hyperparameter configurations are described in §d. additionally, a comparison between ed and do architectures (stage 1) is provided in §b.1; theoretical and empirical complexity analyses are provided in §b.2. ## results the complete results are provided in §e. regarding task performance, our experiments demonstrate that ed models consistently outperform do models across all datasets and models, regardless of their size. presumably, a better inductive bias is injected by applying self-attention (and not autoregressive-attention) to the conditioned input sequence. this finding is particularly relevant for nlp practitioners who aim to develop a specialized in-house model for a specific nlg task. we hence continue to the kd stages only with ed models (t5 and bart). s2: it is better to prune layers from the decoder. in stage 2, we examine whether it is better to prune encoder or decoder layers. to this end, we prune bart-6:6 and report the results at the bottom of second, bart-6:2 outperforms bart-2:6 in every task metric (and dataset), being competitive to bart-6:6. moreover, for tasks with long inputs (e.g., summarization or question generation, s3: use logits kd as the main training objective. in stage 3 we compare different kd objectives. as seen in ## conclusion in this paper, we present a general kd recipe for nlg. to this end, we conduct a systematic study on various tasks and evaluate the impact of different modeling decisions on computational and task performance of distilled models. our results suggest that using ed models as students, pruning decoder layers, combining logits kd and pts via sampling and joint-teaching achieve high compression rates while maintaining competitive performance. nevertheless, our recipe is based on average performance and may depend on the task, model, or setup. the teacher-student performance gap that still exists demonstrate the need for further research. for example, high-temperature pts seem to be less effective for bart, and further exploration of different hyperparameters or methods for increasing pt diversity may be necessary. integrating a smart selection of training examples or pts @xcite , refining joint-teaching with curriculum learning or scheduling @xcite are some future research directions.