Autoregressive Knowledge Distillation through Imitation Learning (2020)

Abstract: The performance of autoregressive models on natural language generation tasks has dramatically improved due to the adoption of deep, self-attentive architectures. However, these gains have come at the cost of hindering inference speed, making state-of-the-art models cumbersome to deploy in real-world, time-sensitive settings. We develop a compression technique for autoregressive models that is driven by an imitation learning perspective on knowledge distillation. The algorithm is designed to address the exposure bias problem. On prototypical language generation tasks such as translation and summarization, our method consistently outperforms other distillation algorithms, such as sequence-level knowledge distillation. Student models trained with our method attain 1.4 to 4.8 BLEU/ROUGE points higher than those trained from scratch, while increasing inference speed by up to 14 times in comparison to the teacher model.

## introduction autoregressive models are ubiquitous in natural language processing. due to the sequential nature of text generation, they are often the tool of choice for tackling sequence-to-sequence problems such as translation @xcite , summarization @xcite , and dialogue @xcite . furthermore, they form the backbone of several successful generative pre-training architectures @xcite @xcite . two recent trends have made autoregressive models cumbersome to deploy in real-world, natural language generation (nlg) applications. first, state-of-the-art models have grown larger and larger, amounting to hundreds of millions and even billions of parameters @xcite @xcite . the increase in size and depth dramatically slows down inference speed. second, the architecture of choice for autoregressive models seems to have shifted from the recurrent neural network (rnn) @xcite to the transformer @xcite . though the transformer's self-attention mechanism improves performance, it also increases the computational complexity of the step-by-step generation algorithms that are used at test time. thus, both of these trends have contributed to significantly increasing inference time costs, especially on cpus and low-resource devices, hindering their use in production systems. knowledge distillation (kd) @xcite is one popular method for model compression. it transfers the information learned by a large, pretrained teacher to a smaller, untrained student. in comparison to other methods such as weight pruning and quantization, kd allows the compressed model's architecture to significantly differ from that of the original teacher. this feature enables models trained with kd to achieve high performance while meeting particular inference requirements (e.g. memory, speed, etc.). sequence-level knowledge distillation (seqkd), proposed by @xcite , is the dominant technique for autoregressive kd in the current nlg literature, especially for machine translation @xcite @xcite . this method trains a student model using a modified dataset generated by the teacher model and the standard negative log-likelihood objective. while seqkd is simple and efficient, we argue that it does not take advantage of the teacher's full potential. training the student model with a static dataset leads to the exposure bias problem. during training, the student model learns to predict the next token given previous tokens provided by the data. however, at inference time, the student generates the entire sequence from scratch by repeatedly using its own outputs as context for subsequent steps. this training-inference inconsistency causes a decrease in generation quality. alternatively, we propose that the student can leverage the teacher in a dynamic fashion during the learning process. we devise a new compression algorithm for autoregressive models called imitation-based knowledge distillation (imitkd). it is inspired by an imitation learning (il) perspective on the autoregressive distillation problem. our algorithm trains a student model within an il framework by treating the teacher as an oracle, and allows the student to explore its own generation during training. the teacher corrects the student's generation at every time step, thereby guiding the student in learning how to generate. experimental results in translation and summarization show that imitkd is especially suitable for compressing deep transformer models that achieve high performance into shallow rnns that generate up to 14 times faster at inference time. our method consistently outperforms other distillation algorithms (such as word-level kd and sequencelevel kd), and yields student models that beat models trained without a teacher by 1.4 to 4.8 points on generation metrics such as bleu and rouge. ## imitation-based distillation algorithm in this section, we present our il-based algorithm for autoregressive distillation. we begin by describing the key design principles and why we expect them to work well. then, we elaborate on the algorithm's implementation in detail. ## related work the distillation problem for autoregressive models was first tackled by @xcite , who introduced sequence-level knowledge distillation for neural machine translation. subsequent works have used seqkd for non-autoregressive translation models @xcite , lowresource settings @xcite , and ensemble distillation with multiple teachers @xcite . @xcite proposed a behavioral cloning method for distilling autoregressive translation models into non-autoregresssive translation models. in contrast, our method aims to address the learning challenges in autoregressive distillation, such as exposure bias. various methods other than standard supervised learning have been explored for training generative models of language. mixer @xcite and beam search optimization @xcite also perform generation during training, but use sequence-level metrics (e.g. bleu score) as training supervision. simlarly, searnn @xcite trains rnns to iteratively generate sequences with beam search to compute the local loss of a single action during the decoding process. scheduled sampling @xcite and its extensions @xcite alleviate exposure bias by replacing some words in the true context with the model's prediction. however, without a dynamic queryable oracle, these methods face the challenge of properly defining the training signal when the generated sequence no longer exists in the static training data. for example, directly reusing the tokens in the static dataset as the target next token leads to an inconsistent training procedure @xcite . in contrast to these methods, distillation can fully leverage the teacher oracle, allowing us to design a simple and efficient imitation learning algorithm. ## experimental setup we test our autoregressive distillation method and all baselines on three language generation tasks -iwslt 2014 @xmath0, wmt 2016 @xmath1, and cnn/dailymail abstractive news summarization. ## conclusion in this work, we developed a new knowledge distillation technique inspired by imitation learning for compressing large and cumbersome autoregressive models into smaller and faster counterparts. we demonstrated the empirical success of our method over popular baselines on several natural language generation tasks. we are excited about several possible avenues for future work. one branch of ideas involves incorporating more advanced il algorithms beyond dagger, such as lols @xcite , to further improve the distillation process. another possibility is to design imitation-based fine-tuning analogs to the seqinter method. finally, although our experiments in this paper focused on sequenceto-sequence settings, we are interested in exploring the use of imitkd for compressing large language models aimed at transfer learning.