X treme D istil: Multi-stage Distillation for Massive Multilingual Models (2020)

Abstract: Deep and large pre-trained language models are the state-of-the-art for various natural language 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 focus on multilingual Named Entity Recognition (NER). In particular, we study several distillation strategies and propose a stage-wise optimization scheme leveraging teacher internal representations, that is agnostic of teacher architecture, and show that it outperforms strategies employed in prior works. Additionally, we investigate the role of several factors like the amount of unlabeled data, annotation resources, model architecture and inference latency 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 latency for batch inference while retaining 95% of its F1-score for NER over 41 languages.

## introduction motivation: pre-trained language models have shown state-of-the-art performance for various natural language processing applications like text classification, named entity recognition and questionanswering. a significant challenge facing practitioners is how to deploy these huge models in practice. for instance, models like bert large @xcite , gpt 2 @xcite , megatron @xcite and t5 @xcite have 340m , 1.5b, 8.3b and 11b parameters 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 @xcite 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 @xcite @xcite leverage soft logits from teachers as optimization targets for distilling students, with some notable exceptions from concurrent work. @xcite ; @xcite ; @xcite ; @xcite additionally use internal teacher representations as additional signals. however, these methods are constrained by architectural considerations like embedding dimension in bert and transformer architecture. this makes it difficult to massively compress models (without being able to reduce network width) or adopt alternate architecture. for instance, we observe bilstms as students to be more accurate than transformers for low latency configurations. some concurrent works @xcite ; @xcite adopt pre-training or dual training to distil students of arbitrary architecture. however, pre-training is expensive in terms of time and computational resources. additionally, most of the above works are geared for distilling language models for glue tasks @xcite . there has been some limited exploration of such techniques for sequence tagging tasks like ner @xcite or multilingual tasks @xcite . however, these works also suffer from similar drawbacks as mentioned before. overview of xtremedistil: in this work, we compare distillation strategies used in all the above xtremedistil: multilingual pre-trained model distillation works and propose a new scheme outperforming prior ones. in this, we leverage teacher internal representations to transfer knowledge to the student. however, in contrast to prior work, we are not restricted by the choice of student architecture. this allows representation transfer from transformerbased teacher model to bilstm-based student model with different embedding dimensions and disparate output spaces. we also propose a stagewise optimization scheme to sequentially transfer most general to task-specific information from teacher to student for better distillation. overview of our task: unlike prior works mostly focusing on glue tasks in a single language, we employ our techniques to study distillation for massive multilingual named entity recognition (ner) over 41 languages. prior work on multilingual transfer on the same @xcite (mm-ner) requires knowledge of source and target language whereby they judiciously select pairs for effective transfer resulting in a customized model for each language. in our work, we adopt multilingual bidirectional encoder representations from transformer (mbert) as our teacher and show that it is possible to perform language-agnostic joint ner for all languages with a single model that has a similar performance but massively compressed in contrast to mbert and mmner. the closest one to this work is that of @xcite where mbert is leveraged for multilingual ner. we discuss this in details and use their strategy as a baseline. we show our distillation strategy to be better leading to a higher compression and faster inference. we also investigate several unexplored dimensions of distillation like the impact of unlabeled transfer data and annotation resources, choice of multilingual word embeddings, architectural variations and inference latency. our techniques obtain massive compression of teacher models like mbert by upto 35x in terms of parameters and 51x in terms of latency for batch inference while retaining 95% of its performance for massive multilingual ner, and matching or outperforming it for classification tasks. overall, our work makes the following contributions: ## models the student: the input to the model are edimensional word embeddings for each token. to capture sequential information in the sentence, we use a single layer bidirectional long short term memory network (bilstm). given a sequence of k tokens, a bilstm computes a set of k vectors h(x k ) = [ ---→ h(x k ); ← --h(x k )] as the concatenation of the states generated by a forward ( ---→ h(x k )) and backward lstm ( ← --h(x k )). assuming the number of hidden units in the lstm to be h, each hidden state h(x k ) is of dimension 2h. probability distribution for the token label at timestep k is given by: we train the student model end-to-end minimizing the above cross-entropy loss over labeled data. the teacher: pre-trained language models like elmo @xcite , bert @xcite and gpt @xcite have shown state-of-the-art performance for several tasks. we adopt bert as the teacher -specifically, the multilingual version of bert (mbert) with 179m m parameters trained over 104 languages with the largest wikipedias. mbert does not use any markers to distinguish languages during pre-training and learns a single language-agnostic model trained via masked language modeling over wikipedia articles from all languages. tokenization: similar to mbert, we use word-piece tokenization with 110k shared wordpiece vocabulary. we preserve casing, remove accents, split on punctuations and whitespace. fine-tuning the teacher: the pre-trained language models are trained for general language modeling objectives. in order to adapt them for the given task, the teacher is fine-tuned end-to-end with task-specific labeled data d l to learn parameters θt using cross-entropy loss as in equation 2 . ## distillation features teacher fine-tuning gives us access to task-specific representations for distilling the student. to this end, we use different kinds of teacher information. ## training we want to optimize the loss functions for representation l rl , logits l ll and cross-entropy l ce . these optimizations can be scheduled differently to obtain different training regimens as follows. ## summary teacher hidden representation and distillation schedule: internal teacher representations help in distillation, although a naive combination hurts the student model. we show that a distillation schedule with stagewise optimization, gradual unfreezing with a cosine learning rate scheduler (d4.1 + d4.2 in