Autoregressive Knowledge Distillation through Imitation Learning

Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, pages 6121–6133, November 16–20, 2020. c⃝2020 Association for Computational Linguistics 6121 Autoregressive Knowledge Distillation through Imitation Learning Alexander Lin1 ASAPP, Inc. New York, NY, USA Jeremy Wohlwend2 ASAPP, Inc. New York, NY, USA Howard Chen3 ASAPP, Inc. New York, NY, USA Tao Lei4 ASAPP, Inc. New York, NY, USA 1alexanderlin01@g.harvard.edu 2jwohlwend@csail.mit.edu 3howardchen@cs.princeton.edu 4tao@asapp.com Abstract The performance of autoregressive models on natural language generation tasks has dramat- ically improved due to the adoption of deep, self-attentive architectures. However, these gains have come at the cost of hindering in- ference speed, making state-of-the-art mod- els 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 de- signed to address the exposure bias problem. On prototypical language generation tasks such as translation and summarization, our method consistently outperforms other distilla- tion algorithms, such as sequence-level knowl- edge 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.1 1 Introduction Autoregressive models are ubiquitous in natural lan- guage 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 (Sutskever et al., 2014), summarization (Rush et al., 2015), and dialogue (Eric and Man- ning, 2017). Furthermore, they form the backbone of several successful generative pre-training archi- tectures (Howard and Ruder, 2018; Peters et al., 2018; Radford et al., 2019; Dai et al., 2019). Two recent trends have made autoregressive models cumbersome to deploy in real-world, natu- ral language generation (NLG) applications. First, state-of-the-art models have grown larger and larger, amounting to hundreds of millions and even 1Our code can be found at https://github.com/ asappresearch/imitkd. billions of parameters (Dong et al., 2019; Liu and Lapata, 2019; Raffel et al., 2019). The increase in size and depth dramatically slows down infer- ence speed. Second, the architecture of choice for autoregressive models seems to have shifted from the recurrent neural network (RNN) (Bahdanau et al., 2014; Luong et al., 2015) to the Transformer (Vaswani et al., 2017). 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 con- tributed to significantly increasing inference time costs, especially on CPUs and low-resource de- vices, hindering their use in production systems. Knowledge distillation (KD) (Buciluˇa et al., 2006; Hinton et al., 2015) 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 meth- ods such as weight pruning and quantization, KD allows the compressed model’s architecture to sig- nificantly differ from that of the original teacher. This feature enables models trained with KD to achieve high …

Content truncated to 500 words