XtremeDistil: Multi-stage Distillation for Massive Multilingual Models

Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2221–2234 July 5 - 10, 2020. c⃝2020 Association for Computational Linguistics 2221 XtremeDistil: Multi-stage Distillation for Massive Multilingual Models Subhabrata Mukherjee Microsoft Research AI Redmond, WA submukhe@microsoft.com Ahmed Hassan Awadallah Microsoft Research AI Redmond, WA hassanam@microsoft.com Abstract Deep and large pre-trained language models are the state-of-the-art for various natural lan- guage processing tasks. However, the huge size of these models could be a deterrent to using them in practice. Some recent works use knowledge distillation to compress these huge models into shallow ones. In this work we study knowledge distillation with a fo- cus on multilingual Named Entity Recognition (NER). In particular, we study several distil- lation strategies and propose a stage-wise op- timization scheme leveraging teacher internal representations, that is agnostic of teacher ar- chitecture, and show that it outperforms strate- gies employed in prior works. Additionally, we investigate the role of several factors like the amount of unlabeled data, annotation re- sources, model architecture and inference la- tency to name a few. We show that our approach leads to massive compression of teacher models like mBERT by upto 35x in terms of parameters and 51x in terms of la- tency for batch inference while retaining 95% of its F1-score for NER over 41 languages. 1 Introduction Motivation: Pre-trained language models have shown state-of-the-art performance for various nat- ural language processing applications like text clas- sification, named entity recognition and question- answering. A significant challenge facing practi- tioners is how to deploy these huge models in prac- tice. For instance, models like BERT Large (Devlin et al., 2019), GPT 2 (Radford et al., 2019), Mega- tron (Shoeybi et al., 2019) and T5 (Raffel et al., 2019) have 340M, 1.5B, 8.3B and 11B parame- ters respectively. Although these models are trained offline, during prediction we need to traverse the deep neural network architecture stack involving a large number of parameters. This significantly increases latency and memory requirements. Knowledge distillation (Hinton et al., 2015; Ba and Caruana, 2014) earlier used in computer vision provides one of the techniques to compress huge neural networks into smaller ones. In this, shallow models (called students) are trained to mimic the output of huge models (called teachers) based on a transfer set. Similar approaches have been recently adopted for language model distillation. Limitations of existing work: Recent works (Liu et al., 2019; Zhu et al., 2019; Tang et al., 2019; Turc et al., 2019) leverage soft logits from teachers as op- timization targets for distilling students, with some notable exceptions from concurrent work. Sun et al. (2019); Sanh (2019); Aguilar et al. (2019); Zhao et al. (2019) additionally use internal teacher rep- resentations as additional signals. However, these methods are constrained by architectural considera- tions like embedding dimension in BERT and trans- former architecture. This makes it difficult to mas- sively compress models (without being able to re- duce network width) or adopt alternate architecture. For instance, we observe BiLSTMS as students …

Content truncated to 500 words