SlowGuess/ACL_2020_OCR huggingface dataset:
Subhabrata Mukherjee
submukhe@microsoft.com
Ahmed Hassan Awadallah
hassanam@microsoft.com
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  $F_{1}$ -score for NER over 41 languages.
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 question-answering. A significant challenge facing practitioners is how to deploy these huge models in practice. For instance, models like BERT Large (Devlin et al., 2019), GPT 2 (Radford et al., 2019), Megaatron (Shoeybi et al., 2019) and T5 (Raffel et al., 2019) 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 (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 optimization 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 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 (Turc et al., 2019); (Zhao et al., 2019) 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 (Wang et al., 2018). There has been some limited exploration of such techniques for sequence tagging tasks like NER (Izsak et al., 2019; Shi et al., 2019) or multilingual tasks (Tsai et al., 2019). 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
Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2221-2234 July 5 - 10, 2020. ©2020 Association for Computational Linguistics
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 Transformer-based 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 (Rahimi et al., 2019) (MMNER) 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 (Tsai et al., 2019) 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:
- Method: We propose a distillation method leveraging internal representations and parameter projection that is agnostic of teacher architecture.
- Inference: To learn model parameters, we propose stage wise optimization schedule with gradual unfreezing outperforming prior schemes.
- Experiments: We perform distillation for multilingual NER on 41 languages with massive compression and comparable performance to huge models<sup>1</sup>. We also perform classification experiments on four datasets where our compressed models perform at par with significantly larger teachers.
- Study: We study the influence of several factors on distillation like the availability of annotation resources for different languages, model architecture, quality of multilingual word embeddings, memory footprint and inference latency.
Problem Statement: Consider a sequence  $x = \langle x_{k} \rangle$  with  $K$  tokens and  $y = \langle y_{k} \rangle$  as the corresponding labels. Consider  $D_{l} = \{\langle x_{k,l} \rangle, \langle y_{k,l} \rangle\}$  to be a set of  $n$  labeled instances with  $X = \{\langle x_{k,l} \rangle\}$  denoting the instances and  $Y = \{\langle y_{k,l} \rangle\}$  the corresponding labels. Consider  $D_{u} = \{\langle x_{k,u} \rangle\}$  to be a transfer set of  $N$  unlabeled instances from the same domain where  $n \ll N$ . Given a teacher  $T(\theta^{t})$ , we want to train a student  $S(\theta^{s})$  with  $\theta$  being trainable parameters such that  $|\theta^{s}| \ll |\theta^{t}|$  and the student is comparable in performance to the teacher based on some evaluation metric. In the following section, the superscript 't' always represents the teacher and 's' denotes the student.
Model compression and knowledge distillation: Prior works in the vision community dealing with huge architectures like AlexNet and ResNet have addressed this challenge in two ways. Works in model compression use quantization (Gong et al., 2014), low-precision training and pruning the network, as well as their combination (Han et al., 2016) to reduce the memory footprint. On the other hand, works in knowledge distillation leverage student teacher models. These approaches include using soft logits as targets (Ba and Caruana, 2014), increasing the temperature of the softmax to match that of the teacher (Hinton et al., 2015) as well as using teacher representations (Romero et al., 2015) (refer to (Cheng et al., 2017) for a survey).
Recent and concurrent Works: Liu et al. (2019); Zhu et al. (2019); Clark et al. (2019) leverage ensembling to distil knowledge from several multitask deep neural networks into a single model. Sun et al. (2019); Sanh (2019); Aguilar et al. (2019) train student models leveraging architectural knowledge
1Code and resources available at: https://aka.ms/XtremeDistil
of the teacher models which adds architectural constraints (e.g., embedding dimension) on the student. In order to address this shortcoming, more recent works combine task-specific distillation with pre-training the student model with arbitrary embedding dimension but still relying on transformer architectures (Turc et al., 2019); (Jiao et al., 2019); (Zhao et al., 2019).
Izsak et al. (2019); Shi et al. (2019) extend these for sequence tagging for Part-of-Speech (POS) tagging and Named Entity Recognition (NER) in English. The one closest to our work Tsai et al. (2019) extends the above for multilingual NER.
Most of these works rely on general corpora for pre-training and task-specific labeled data for distillation. To harness additional knowledge, (Turc et al., 2019) leverage task-specific unlabeled data. (Tang et al., 2019; Jiao et al., 2019) use rule-and embedding-based data augmentation.
The Student: The input to the model are  $E$ -dimensional 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}) = \overleftarrow{[h(x_{k})};\overleftarrow{h(x_{k})]}$  as the concatenation of the states generated by a forward  $(\overleftarrow{h(x_k)})$  and backward LSTM  $(\overleftarrow{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:
$$
p ^ {(s)} (x _ {k}) = \operatorname {s o f t m a x} (h (x _ {k}) \cdot W ^ {s}) \qquad (1)
$$
where  $W^{s} \in R^{2H.C}$  and  $C$  is number of labels.
Consider one-hot encoding of the token labels, such that  $y_{k,l,c} = 1$  for  $y_{k,l} = c$ , and  $y_{k,l,c} = 0$  otherwise for  $c \in C$ . The overall cross-entropy loss computed over each token obtaining a specific label in each sequence is given by:
$$
\mathcal {L} _ {\mathcal {C E}} = - \sum_ {x _ {l}, y _ {l} \in D _ {l}} \sum_ {k} \sum_ {c} y _ {k, c, l} \log p _ {c} ^ {(s)} \left(x _ {k, l}\right) (2)
$$
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 (Peters et al., 2018), BERT (Devlin et al., 2019) and GPT (Radford et al., 2018, 2019) have shown state-of-the-art performance for several tasks. We adopt BERT as the teacher - specifically,
the multilingual version of BERT (mBERT) with 179MM 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 WordPiece 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  $\tilde{\theta}^{t}$  using cross-entropy loss as in Equation 2.
4 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.
4.1 Teacher Logits
Logits as logarithms of predicted probabilities provide a better view of the teacher by emphasizing on the different relationships learned by it across different instances. Consider  $p^t(x_k)$  to be the classification probability of token  $x_k$  as generated by the fine-tuned teacher with  $\text{logit}(p^t(x_k))$  representing the corresponding logits. Our objective is to train a student model with these logits as targets. Given the hidden state representation  $h(x_k)$  for token  $x_k$ , we can obtain the corresponding classification score (since targets are logits) as:
$$
r ^ {s} \left(x _ {k}\right) = W ^ {r} \cdot h \left(x _ {k}\right) + b ^ {r} \tag {3}
$$
where  $W^{r} \in R^{C \cdot 2H}$  and  $b^{r} \in R^{C}$  are trainable parameters and  $C$  is the number of classes. We want to train the student neural network end-to-end by minimizing the element-wise mean-squared error between the classification scores given by the student and the target logits from the teacher as:
$$
\mathcal {L} _ {\mathcal {L L}} = \frac {1}{2} \sum_ {x _ {u} \in D _ {u}} \sum_ {k} \left| \left| r ^ {s} \left(x _ {k, u}\right) - \operatorname {l o g i t} \left(p ^ {t} \left(x _ {k, u}; \tilde {\theta} _ {t}\right)\right) \right| \right| ^ {2} \tag {4}
$$
4.2 Internal Teacher Representations
Hidden representations: Recent works (Sun et al., 2019; Romero et al., 2015) have shown the
hidden state information from the teacher to be helpful as a hint-based guidance for the student. Given a large collection of task-specific unlabeled data, we can transfer the teacher's knowledge to the student via its hidden representations. However, this poses a challenge in our setting as the teacher and student models have different architectures with disparate output spaces.
Consider  $h^s (x_k)$  and  $z_l^t (x_k;\tilde{\theta}_t)$  to be the representations generated by the student and the  $l^{th}$  deep layer of the fine-tuned teacher respectively for a token  $x_{k}$ . Consider  $x_{u}\in D_{u}$  to be the set of unlabeled instances. We will later discuss the choice of the teacher layer  $l$  and its impact on distillation.
Projection: To make all output spaces compatible, we perform a non-linear projection of the parameters in student representation  $h^s$  to have same shape as teacher representation  $z_l^t$  for each token  $x_k$ :
$$
\tilde {z} ^ {s} (x _ {k}) = G e l u (W ^ {f} \cdot h ^ {s} (x _ {k}) + b ^ {f}) \qquad (5)
$$
where  $W^{f} \in R^{|z_{l}^{t}| \cdot 2H}$  is the projection matrix,  $b^{f} \in R^{|z_{l}^{t}|}$  is the bias, and  $Gelu$  (Gaussian Error Linear Unit) (Hendrycks and Gimpel, 2016) is the non-linear projection function.  $|z_{l}^{t}|$  represents the embedding dimension of the teacher. This transformation aligns the output spaces of the student and teacher and allows us to accommodate arbitrary student architecture. Also note that the projections (and therefore the parameters) are shared across tokens at different timepoints.
The projection parameters are learned by minimizing the  $KL$ -divergence (KLD) between the student and the  $l^{th}$  layer teacher representations:
$$
\mathcal {L} _ {\mathcal {R L}} = \sum_ {x _ {u} \in D _ {u}} \sum_ {k} \operatorname {K L D} \left(\tilde {z} ^ {s} \left(x _ {k, u}\right), z _ {l} ^ {t} \left(x _ {k, u}; \tilde {\theta} _ {t}\right)\right) \tag {6}
$$
Multilingual word embeddings: A large number of parameters reside in the word embeddings. For mBERT a shared multilingual WordPiece vocabulary of  $V = 110K$  tokens and embedding dimension of  $D = 768$  leads to 92MM parameters. To have massive compression, we cannot directly incorporate mBERT embeddings in our model. Since we use the same WordPiece vocabulary, we are likely to benefit more from these embeddings than from Glove (Pennington et al., 2014) or FastText (Bojanowski et al., 2016).
We use a dimensionality reduction algorithm like Singular Value Decomposition (SVD) to project the mBERT word embeddings to a lower dimensional space. Given mBERT word embedding ma
trix of dimension  $V \times D$ , SVD finds the best  $E$ -dimensional representation that minimizes sum of squares of the projections (of rows) to the subspace.
5 Training
We want to optimize the loss functions for representation  $\mathcal{L}_{\mathcal{R},\mathcal{L}}$ , logits  $\mathcal{L}_{\mathcal{L},\mathcal{L}}$  and cross-entropy  $\mathcal{L}_{\mathcal{C},\mathcal{E}}$ . These optimizations can be scheduled differently to obtain different training regimens as follows.
5.1 Joint Optimization
In this, we optimize the following losses jointly:
$$
\frac {1}{| D _ {l} |} \sum_ {\left\{x _ {l}, y _ {l} \right\} \in D _ {l}} \alpha \cdot \mathcal {L} _ {\mathcal {C}} \varepsilon \left(x _ {l}, y _ {l}\right) +
$$
$$
\frac {1}{| D _ {u} |} \sum_ {\left\{x _ {u}, y _ {u} \right\} \in D _ {u}} \left(\beta \cdot \mathcal {L} _ {\mathcal {R L}} \left(x _ {u}, y _ {u}\right) + \gamma \cdot \mathcal {L} _ {\mathcal {L L}} \left(x _ {u}, y _ {u}\right)\right) \tag {7}
$$
where  $\alpha, \beta$  and  $\gamma$  weigh the contribution of different losses. A high value of  $\alpha$  makes the student focus more on easy targets; whereas a high value of  $\gamma$  leads focus to the difficult ones. The above loss is computed over two different task-specific data segments. The first part involves cross-entropy loss over labeled data, whereas the second part involves representation and logit loss over unlabeled data.
5.2 Stage-wise Training
Instead of optimizing all loss functions jointly, we propose a stage-wise scheme to gradually transfer most general to task-specific representations from teacher to student. In this, we first train the student to mimic teacher representations from its  $l^{th}$  layer by optimizing  $\mathcal{R}_{\mathcal{RL}}$  on unlabeled data. The student learns the parameters for word embeddings  $(\theta^w)$ , BiLSTM  $(\theta^b)$  and projections  $\langle W^f, b^f \rangle$ .
In the second stage, we optimize for the cross-entropy  $\mathcal{R}_{\mathcal{CE}}$  and logit loss  $\mathcal{R}_{\mathcal{LL}}$  jointly on both labeled and unlabeled data respectively to learn the corresponding parameters  $W^{s}$  and  $\langle W^r,b^r\rangle$
The above can be further broken down in two stages, where we sequentially optimize logit loss  $\mathcal{R}_{\mathcal{LL}}$  on unlabeled data and then optimize cross-entropy loss  $\mathcal{R}_{\mathcal{CE}}$  on labeled data. Every stage learns parameters conditioned on those learned in previous stage followed by end-to-end fine-tuning.
5.3 Gradual Unfreezing
One potential drawback of end-to-end fine-tuning for stage-wise optimization is 'catastrophic forgetting' (Howard and Ruder, 2018) where the model forgets information learned in earlier stages. To address this, we adopt gradual unfreezing – where we tune the model one layer at a time starting from the configuration at the end of previous stage.
We start from the top layer that contains the most task-specific information and allow the model to configure the task-specific layer first while others remain frozen. The latter layers are gradually unfrozen one by one and the model trained till convergence. Once a layer is unfrozen, it maintains the state. When the last layer (word embeddings) is unfrozen, the entire network is trained end-to-end. The order of this unfreezing scheme (top-to-bottom) is reverse of that in (Howard and Ruder, 2018) and we find this to work better in our setting with the following intuition. At the end of the first stage on optimizing  $\mathcal{R}_{\mathcal{RL}}$ , the student learns to generate representations similar to that of the  $l^{th}$  layer of the teacher. Now, we need to add only a few task-specific parameters  $(\langle W^r, b^r \rangle)$  to optimize for logit loss  $\mathcal{R}_{\mathcal{LL}}$  with all others frozen. Next, we gradually give the student more flexibility to optimize for task-specific loss by tuning the layers below where the number of parameters increases with depth  $(|\langle W^r, b^r \rangle| \ll |\theta_b| \ll |\theta_w|)$ .
We tune each layer for  $n$  epochs and restore model to the best configuration based on validation
Dataset	Labels	Train	Test	Unlabeled
NER
Wikiann-41	11	705K	329K	7.2MM
Classification
IMDB	2	25K	25K	50K
DBPedia	14	560K	70K	-
AG News	4	120K	7.6K	-
Elec	2	25K	25K	200K
Table 1: Full dataset summary.
Work	PT	TA	Distil.
Sanh (2019)	Y	Y	D1
Turc et al. (2019)	Y	N	D1
Liu et al. (2019); Zhu et al. (2019); Shi et al. (2019); Tsai et al. (2019); Tang et al. (2019); Izsak et al. (2019); Clark et al. (2019)	N	N	D1
Sun et al. (2019)	N	Y	D2
Jiao et al. (2019)	N	N	D2
Zhao et al. (2019)	Y	N	D2
XtremeDistil (ours)	N	N	D4
Table 2: Different distillation strategies. D1 leverages soft logits with hard labels. D2 uses representation loss. PT denotes pre-training with language modeling. TA depicts students constrained by teacher architecture.
loss on a held-out set. Therefore, the model retains best possible performance from any iteration. Algorithm 1 shows overall processing scheme.
Dataset Description: We evaluate our model XtremeDistil for multilingual NER on 41 languages and same setting as in (Rahimi et al., 2019). This data is derived from WikiAnn NER corpus (Pan et al., 2017) and partitioned into training, development and test sets. All NER results are reported in this test set for a fair comparison between existing works. We report the average  $F_{1}$ -score ( $\mu$ ) and standard deviation  $\sigma$  between scores across 41 languages for phrase-level evaluation. Refer to Figure 2 for language codes and corresponding distribution of training labels. We also perform experiments with data from four other domains (refer to Table 1): IMDB (Maas et al., 2011), SST-2 (Socher et al., 2013) and Elec (McAuley and Leskovec, 2013) for sentiment analysis for movie and electronics product reviews, DbPedia (Zhang et al., 2015) and Ag News (Zhang et al., 2015) for topic classification of Wikipedia and news articles. NER Tags: The NER corpus uses IOB2 tagging strategy with entities like LOC, ORG and PER. Following mBERT, we do not use language markers and share these tags across all languages. We
Strategy	Features	Transfer = 0.7MM	Transfer = 1.4MM	Transfer = 7.2MM
D0	Labels per lang.	71.26 (6.2)	-	-
D0-S	Labels across all lang.	81.44 (5.3)	-	-
D1	Labels and Logits	82.74 (5.1)	84.52 (4.8)	85.94 (4.8)
D2	Labels, Logits and Repr.	82.38 (5.2)	83.78 (4.9)	85.87 (4.9)
D3.1	(S1) Repr. (S2) Labels and Logits	83.10 (5.0)	84.38 (5.1)	86.35 (4.9)
D3.2	+ Gradual unfreezing	86.77 (4.3)	87.79 (4.0)	88.26 (4.3)
D4.1	(S1) Repr. (S2) Logits (S3) Labels	84.82 (4.7)	87.07 (4.2)	87.87 (4.1)
D4.2	+ Gradual unfreezing	87.10 (4.2)	88.64 (3.8)	88.52 (4.1)
use additional syntactic markers like {CLS, SEP, PAD} and 'X' for marking segmented wordpieces contributing a total of 11 tags (with shared 'O').
6.1 Evaluating Distillation Strategies
Baselines: A trivial baseline (D0) is to learn models one per language using only corresponding labels for learning. This can be improved by merging all instances and sharing information across all languages (D0-S). Most of the concurrent and recent works (refer to Table 2 for an overview) leverage logits as optimization targets for distillation (D1). A few exceptions also use teacher internal representations along with soft logits (D2). For our model we consider multi-stage distillation, where we first optimize representation loss followed by jointly optimizing logit and cross-entropy loss (D3.1) and further improving it by gradual unfreezing of neural network layers (D3.2). Finally, we optimize the loss functions sequentially in three stages (D4.1) and improve it further by unfreezing mechanism (D4.2). We further compare all strategies while varying the amount of unlabeled transfer data for distillation (hyper-parameter settings in Appendix).