SlowGuess/NAACL_2022_OCR huggingface dataset:
Simiao Zuo $^{\ddagger}$ , Qingru Zhang $^{\ddagger}$ , Chen Liang $^{\ddagger}$ , Pengcheng He $^{\diamond}$ , Tuo Zhao $^{\ddagger}$  and Weizhu Chen $^{\diamond}$
$^{\ddagger}$ Georgia Institute of Technology  ${}^{\circ}$ Microsoft
{simiaozuo,qzhang441,cliang73,tourzhao}@gatech.edu
{Pengcheng.H, wzchen}@microsoft.com
Pre-trained language models have demonstrated superior performance in various natural language processing tasks. However, these models usually contain hundreds of millions of parameters, which limits their practicality because of latency requirements in real-world applications. Existing methods train small compressed models via knowledge distillation. However, performance of these small models drops significantly compared with the pretrained models due to their reduced model capacity. We propose MoEBERT, which uses a Mixture-of-Experts structure to increase model capacity and inference speed. We initialize MoEBERT by adapting the feed-forward neural networks in a pre-trained model into multiple experts. As such, representation power of the pre-trained model is largely retained. During inference, only one of the experts is activated, such that speed can be improved. We also propose a layer-wise distillation method to train MoEBERT. We validate the efficiency and effectiveness of MoEBERT on natural language understanding and question answering tasks. Results show that the proposed method outperforms existing task-specific distillation algorithms. For example, our method outperforms previous approaches by over  $2\%$  on the MNLI (mismatched) dataset. Our code is publicly available at https://github.com/ SimiaoZuo/MoEBERT.
Pre-trained language models have demonstrated superior performance in various natural language processing tasks, such as natural language understanding (Devlin et al., 2019; Liu et al., 2019; He et al., 2021b) and natural language generation (Radford et al., 2019; Brown et al., 2020). These models can contain billions of parameters, e.g., T5 (Raffel et al., 2019) contains up to 11 billion parameters, and GPT-3 (Brown et al., 2020) consists of up to
175 billion parameters. Their extreme sizes bring challenges in serving the models to real-world applications due to latency requirements.
Model compression through knowledge distillation (Romero et al., 2015; Hinton et al., 2015) is a promising approach that reduces the computational overhead of pre-trained language models while maintaining their superior performance. In knowledge distillation, a large pre-trained language model serves as a teacher, and a smaller student model is trained to mimic the teacher's behavior. Distillation approaches can be categorized into two groups: task-agnostic (Sanh et al., 2019; Jiao et al., 2020; Wang et al., 2020, 2021; Sun et al., 2020a) and task-specific (Turc et al., 2019; Sun et al., 2019; Li et al., 2020; Hou et al., 2020; Sun et al., 2020b; Xu et al., 2020). Task-agnostic distillation pre-trains the student and then fine-tunes it on downstream tasks; while task-specific distillation directly fine-tunes the student after initializing it from a pre-trained model. Note that task-agnostic approaches are often combined with task-specific distillation during fine-tuning for better performance (Jiao et al., 2020). We focus on task-specific distillation in this work.
One major drawback of existing knowledge distillation approaches is the drop in model performance caused by the reduced representation power. That is, because the student model has fewer parameters than the teacher, its model capacity is smaller. For example, the student model in DistilBERT (Sanh et al., 2019) has 66 million parameters, about half the size of the teacher (BERT-base, Devlin et al. 2019). Consequently, performance of DistilBERT drops significantly compared with BERT-base, e.g., over  $2\%$  on MNLI (82.2 v.s. 84.5) and over  $3\%$  on CoLA (54.7 v.s. 51.3).
We resort to the Mixture-of-Experts (MoE, Shazeer et al. 2017) structure to remedy the representation power issue. MoE models can increase model capacity while keeping the inference com
Human Language Technologies, pages 1610 - 1623
putational cost constant. A layer of a MoE model (Shazeer et al., 2017; Lepikhin et al., 2021; Fedus et al., 2021; Yang et al., 2021; Zuo et al., 2021) consists of an attention mechanism and multiple feed-forward neural networks (FFNs) in parallel. Each of the FFNs is called an expert. During training and inference, an input adaptively activates a fixed number of experts (usually one or two). In this way, the computational cost of a MoE model remains constant during inference, regardless of the total number of experts. Such a property facilitates compression without reducing model capacity.
However, MoE models are difficult to train-from-scratch and usually require a significant amount of parameters, e.g., 7.4 billion parameters for Switchbase (Fedus et al., 2021). We propose MoEBERT, which incorporates the MoE structure into pretrained language models for fine-tuning. Our model can speedup inference while retaining the representation power of the pre-trained language model. Specifically, we incorporate the expert structure by adapting the FFNs in a pre-trained model into multiple experts. For example, the hidden dimension of the FFN is 3072 in BERT-base (Devlin et al., 2019), and we adapt it into 4 experts, each has a hidden dimension 768. In this way, the amount of effective parameters (i.e., parameters involved in computing the representation of an input) is cut by half, and we obtain a  $\times 2$  speedup. We remark that MoEBERT utilizes more parameters of the pretrained model than existing approaches, such that it has greater representation power.
To adapt the FFNs into experts, we propose an importance-based method. Empirically, there are some neurons in the FFNs that contribute more to the model performance than the other ones. That is, removing the important neurons causes significant performance drop. Such a property can be quantified by the importance score (Molchanov et al., 2019; Xiao et al., 2019; Liang et al., 2021). When initializing MoEBERT, we share the most important neurons (i.e., the ones with the highest scores) among the experts, and the other neurons are distributed evenly. This strategy has two advantages: first, the shared neurons preserve performance of the pre-trained model; second, the non-shared neurons promote diversity among experts, which further boost model performance. After initialization, MoEBERT is trained using a layer-wise task-specific distillation algorithm.
We demonstrate efficiency and effectiveness of
MoEBERT on natural language understanding and question answering tasks. On the GLUE (Wang et al., 2019) benchmark, our method significantly outperforms existing distillation algorithms. For example, MoEBERT exceeds performance of state-of-the-art task-specific distillation approaches by over  $2\%$  on the MNLI (mismatched) dataset. For question answering, MoEBERT increases F1 by 2.6 on SQuAD v1.1 (Rajpurkar et al., 2016) and 7.0 on SQuAD v2.0 (Rajpurkar et al., 2018) compared with existing algorithms.
The rest of the paper is organized as follows: we introduce background and related works in Section 2; we describe MoEBERT in Section 3; experimental results are provided in Section 4; and Section 5 concludes the paper.
2.1 Backbone: Transformer
The Transformer (Vaswani et al., 2017) backbone has been widely adopted in pre-trained language models. The model contains several identically-constructed Transformer layers. Each layer has a multi-head self-attention mechanism and a two-layer feed-forward neural network (FFN).
Suppose the output of the attention mechanism is A. Then, the FFN is defined as:
$$
\mathbf {H} = \sigma \left(\mathbf {A} \mathbf {W} _ {1} + \mathbf {b} _ {1}\right), \mathbf {X} ^ {\ell} = \mathbf {W} _ {2} \mathbf {H} + \mathbf {b} _ {2}, \tag {1}
$$
where  $\mathbf{W}_1\in \mathbb{R}^{d\times d_h}$ ,  $\mathbf{W}_2\in \mathbb{R}^{d_h\times d}$ ,  $\mathbf{b}_1\in \mathbb{R}^{d_h}$  and  $\mathbf{b}_2\in \mathbb{R}^d$  are weights of the FFN, and  $\sigma$  is the activation function. Here  $d$  denotes the embedding dimension, and  $d_h$  denotes the hidden dimension of the FFN.
2.2 Mixture-of-Experts Models
Mixture-of-Experts models consist of multiple expert layers, which are similar to the Transformer layers. Each of these layers contains a self-attention mechanism and multiple FFNs (Eq. 1) in parallel, where each FFN is called an expert.
Let  $\{E_i\}_{i=1}^N$  denote the experts, and  $N$  denotes the total number of experts. Similar to Eq. 1, the experts in layer  $\ell$  take the attention output  $\mathbf{A}$  as the input. For each  $\mathbf{a}_t$  (the  $t$ -th row of  $\mathbf{A}$ ) that corresponds to an input token, the corresponding output  $\mathbf{x}_t^\ell$  of layer  $\ell$  is
$$
\mathbf {x} _ {t} ^ {\ell} = \sum_ {i \in \mathcal {T}} p _ {i} (\mathbf {a} _ {t}) E _ {i} (\mathbf {a} _ {t}). \tag {2}
$$
Here,  $\mathcal{T} \subset \{1 \cdots N\}$  is the activated set of experts with  $|\mathcal{T}| = K$ , and  $p_i$  is the weight of expert  $E_i$ .
Different approaches have been proposed to construct  $\mathcal{T}$  and compute  $p_i$ . For example, Shazeer et al. (2017) take
$$
p _ {i} \left(\mathbf {a} _ {t}\right) = \left[ \operatorname {s o f t m a x} \left(\mathbf {a} _ {t} \mathbf {W} _ {g}\right) \right] _ {i}, \tag {3}
$$
where  $\mathbf{W}_g$  is a weight matrix. Consequently,  $\mathcal{T}$  is constructed as the experts that yield top- $K$  largest  $p_i$ . However, such an approach suffers from load imbalance, i.e.,  $\mathbf{W}_g$  collapses such that nearly all the inputs are routed to the same expert. Existing works adopt various ad-hoc heuristics to mitigate this issue, e.g., adding Gaussian noise to Eq. 3 (Shazeer et al., 2017), limiting the maximum number of inputs that can be routed to an expert (Lepikhin et al., 2021), imposing a load balancing loss (Lepikhin et al., 2021; Fedus et al., 2021), and using linear assignment (Lewis et al., 2021). In contrast, Roller et al. 2021 completely remove the gate and pre-assign tokens to experts using hash functions, in which case we can take  $p_i = 1 / K$ .
In Eq. 2, a token only activates  $K$  instead of  $N$  experts, and usually  $K \ll N$ , e.g.,  $K = 2$  and  $N = 2048$  in GShard (Lepikhin et al., 2021). As such, the number of FLOPs for one forward pass does not scale with the number of experts. Such a property paves the way for increasing inference speed of a pre-trained model without decreasing the model capacity, i.e., we can adapt the FFNs in a pretrained model into several smaller components, and only activate one of the components for a specific input token.
2.3 Pre-trained Language Models
Pre-trained language models (Peters et al., 2018; Devlin et al., 2019; Raffel et al., 2019; Liu et al., 2019; Brown et al., 2020; He et al., 2021b,a) have demonstrated superior performance in various natural language processing tasks. These models are trained on an enormous amount of unlabeled data, such that they contain rich semantic information that benefits downstream tasks. Fine-tuning pretrained language models achieves state-of-the-art performance in tasks such that natural language understanding (He et al., 2021a) and natural language generation (Brown et al., 2020).
2.4 Knowledge Distillation
Knowledge distillation (Romero et al., 2015; Hinton et al., 2015) compensates for the performance drop caused by model compression. In knowledge distillation, a small student model mimics the behavior of a large teacher model. For example, DistilBERT (Sanh et al., 2019) uses the teacher's soft
prediction probability to train the student model; TinyBERT (Jiao et al., 2020) aligns the student's layer outputs (including attention outputs and hidden states) with the teacher's; MiniLM (Wang et al., 2020, 2021) utilizes self-attention distillation; and CoDIR (Sun et al., 2020a) proposes to use a contrastive objective such that the student can distinguish positive samples from negative ones according to the teacher's outputs.
There are also heated discussions on the number of layers to distill. For example, Wang et al. (2020, 2021) distill the attention outputs of the last layer; Sun et al. (2019) choose specific layers to distill; and Jiao et al. (2020) use different weights for different transformer layers.
There are two variants of knowledge distillation: task-agnostic (Sanh et al., 2019; Jiao et al., 2020; Wang et al., 2020, 2021; Sun et al., 2020a) and task-specific (Turc et al., 2019; Sun et al., 2019; Li et al., 2020; Hou et al., 2020; Sun et al., 2020b; Xu et al., 2020). The former requires pre-training a small model using knowledge distillation and then fine-tuning on downstream tasks, while the latter directly fine-tunes the small model. Note that task-agnostic approaches are often combined with task-specific distillation for better performance, e.g., TinyBERT (Jiao et al., 2020). In this work, we focus on task-specific distillation.
In this section, we first present an algorithm that adapts a pre-trained language model into a MoE model. Such a structure enables inference speedup by reducing the number of parameters involved in computing an input token's representation. Then, we introduce a layer-wise task-specific distillation method that compensates for the performance drop caused by model compression.
3.1 Importance-Guided Adaptation of Pre-trained Language Models
Adapting the FFNs in a pre-trained language model into multiple experts facilitates inference speedup while retaining model capacity. This is because in a MoE model, only a subset of parameters are used to compute the representation of a given token (Eq. 2). These activated parameters are referred to as effective parameters. For example, by adapting the FFNs in a pre-trained BERT-base (Devlin et al., 2019) (with hidden dimension 3072) model into 4 experts (each has hidden dimension 768), the
Figure 1: Adapting a two-layer FFN into two experts. The blue neuron is the most important one, and is shared between the two experts. The red and green neurons are the second and third important ones, and are assigned to expert one and two, respectively.
number of effective parameters reduces by half, such that we obtain a  $\times 2$  speedup.
Empirically, we find that randomly converting a FFN into experts works poorly (see Figure 3a in the experiments). This is because there are some columns in  $\mathbf{W}_1\in \mathbb{R}^{d\times d_h}$  (correspondingly some rows in  $\mathbf{W}_2$  in Eq. 1) contribute more than the others to model performance.
The importance score (Molchanov et al., 2019; Xiao et al., 2019; Liang et al., 2021), originally introduced in model pruning literature, measures such parameter importance. For a dataset  $\mathcal{D}$  with sample pairs  $\{(x,y)\}$ , the score is defined as
$$
\begin{array}{l} I _ {j} = \sum_ {(x, y) \in \mathcal {D}} \left| \left(\mathbf {w} _ {j} ^ {1}\right) ^ {\top} \nabla_ {\mathbf {w} _ {j} ^ {1}} \mathcal {L} (x, y) \right. \\ + \left(\mathbf {w} _ {j} ^ {2}\right) ^ {\top} \nabla_ {\mathbf {w} _ {j} ^ {2}} \mathcal {L} (x, y) \Bigg |. \tag {4} \\ \end{array}
$$
Here  $\mathbf{w}_j^1\in \mathbb{R}^d$  is the  $j$ -th column of  $\mathbf{W}_1$ ,  $\mathbf{w}_j^2$  is the  $j$ -th row of  $\mathbf{W}_2$ , and  $\mathcal{L}(x,y)$  is the loss.
The importance score in Eq. 4 indicates variation of the loss if we remove the neuron. That is,
$$
\begin{array}{l} \left| \mathcal {L} _ {\mathbf {w}} - \mathcal {L} _ {\mathbf {w} = \mathbf {0}} \right| \approx \left| (\mathbf {w} - \mathbf {0}) ^ {\top} \nabla_ {\mathbf {w}} \mathcal {L} _ {\mathbf {w}} \right| \\ = | \mathbf {w} ^ {\top} \nabla_ {\mathbf {w}} \mathcal {L} _ {\mathbf {w}} |, \\ \end{array}
$$
where  $\mathcal{L}_{\mathbf{w}}$  is the loss with neuron<sup>1</sup> w and  $\mathcal{L}_{\mathbf{w} = 0}$  is the loss without neuron w. Here the approximation is based on the first order Taylor expansion of  $\mathcal{L}_{\mathbf{w}}$  around  $\mathbf{w} = \mathbf{0}$ .
After computing  $I_{j}$  for all the columns, we adapt  $\mathbf{W}_1$  into experts. The columns are sorted in as
cending order according to their importance scores as  $\mathbf{w}_{(1)}^{1}\dots \mathbf{w}_{(d_h)}^{1}$ , where  $\mathbf{w}_{(1)}^{1}$  has the largest  $I_{j}$  and  $\mathbf{w}_{(d_h)}^{1}$  the smallest. Empirically, we find that sharing the most important columns benefits model performance. Based on this finding, suppose we share the top- $s$  columns and we adapt the FFN into  $N$  experts, then expert  $e$  contains columns  $\{\mathbf{w}_{(1)}^{1},\dots ,\mathbf{w}_{(s)}^{1},\mathbf{w}_{(s + e)}^{1},\mathbf{w}_{(s + e + N)}^{1},\dots \}$ . Note that we discard the least important columns to keep the size of each expert as  $\lfloor d / N\rfloor$ . Figure 1 is an illustration of adapting a FFN with 4 neurons in a pre-trained model into two experts.
3.2 Layer-wise Distillation
To remedy the performance drop caused by adapting a pre-trained model to a MoE model, we adopt a layer-wise task-specific distillation algorithm. We use BERT-base (Devlin et al., 2019) as both the student (i.e., the MoE model) and the teacher. We distill both the Transformer layer output  $\mathbf{X}^{\ell}$  (Eq. 2) and the final prediction probability.
For the Transformer layers, the distillation loss is the mean squared error between the teacher's layer output  $\mathbf{X}_{\mathrm{tea}}^{\ell}$  and the student's layer output  $\mathbf{X}^{\ell}$  obtained from Eq. 2.3. Concretely, for an input  $x$ , the Transformer layer distillation loss is
$$
\mathcal {L} _ {\mathrm {t r m}} (x) = \sum_ {\ell = 0} ^ {L} \operatorname {M S E} \left(\mathbf {X} ^ {\ell}, \mathbf {X} _ {\text {t e a}} ^ {\ell}\right), \tag {5}
$$
where  $L$  is the total number of layers. Notice that we include the MSE loss of the embedding layer outputs  $\mathbf{X}^0$  and  $\mathbf{X}_{\mathrm{tea}}^0$ .
Let  $f$  denotes the MoE model and  $f_{\mathrm{tea}}$  the teacher model. We obtain the prediction probability for an input  $x$  as  $p = f(x)$  and  $p_{\mathrm{tea}} = f_{\mathrm{tea}}(x)$ , where  $p$  is the prediction of the MoE model and  $p_{\mathrm{tea}}$  is the prediction of the teacher model. Then the distillation loss for the prediction layer is
$$
\mathcal {L} _ {\text {p r e d}} (x) = \frac {1}{2} \left(\mathrm {K L} \left(p \| p _ {\text {t e a}}\right) + \mathrm {K L} \left(p _ {\text {t e a}} \| p\right)\right), \tag {6}
$$
where KL is the Kullback-Leibler divergence.
The layer-wise distillation loss is the sum of Eq. 5 and Eq. 6, defined as
$$
\mathcal {L} _ {\text {d i s t i l l}} (x) = \mathcal {L} _ {\text {t r m}} (x) + \mathcal {L} _ {\text {p r e d}} (x). \tag {7}
$$
We will discuss variants of Eq. 7 in the experiments.
1 A neuron  $\mathbf{w}$  contains two weights  $\mathbf{w}^1$  and  $\mathbf{w}^2$  as in Eq. 4.
2The other parameters in the FFN:  $\mathbf{W}_2$ $\mathbf{b}_1$  and  $\mathbf{b}_2$  are treated similarly according to  $\{I_j\}$
<sup>3</sup>Note that Eq. 2 computes the layer output of one token  $\mathbf{x}_t^\ell$ , i.e., one row in  $\mathbf{X}^\ell$ .
We employ the random hashing strategy (Roller et al., 2021) to train the experts. That is, each token is pre-assigned to a random expert, and this assignment remains the same during training and inference. We will discuss more about other routing strategies of the MoE model in the experiments.
Given the training dataset  $\mathcal{D}$  and samples  $\{(x,y)\}$ , the training objective is
$$
\mathcal {L} = \sum_ {(x, y) \in \mathcal {D}} \operatorname {C E} (f (x), y) + \lambda_ {\text {d i s t i l l}} \mathcal {L} _ {\text {d i s t i l l}} (x),
$$
where CE is the cross-entropy loss and  $\lambda_{\mathrm{distill}}$  is a hyper-parameter.
In this section, we evaluate the effectiveness and efficiency of the proposed algorithm on natural language understanding and question answering tasks. We implement our algorithm using the Huggingface Transformers $^4$  (Wolf et al., 2019) code-base. All the experiments are conducted on NVIDIA V100 GPUs.
GLUE. We evaluate performance of the proposed method on the General Language Understanding Evaluation (GLUE) benchmark (Wang et al., 2019), which is a collection of nine natural language understanding tasks. The benchmark includes two single-sentence classification tasks: SST-2 (Socher et al., 2013) is a binary classification task that classifies movie reviews to positive or negative, and CoLA (Warstadt et al., 2019) is a linguistic acceptability task. GLUE also contains three similarity and paraphrase tasks: MRPC (Dolan and Brockett, 2005) is a paraphrase detection task; STS-B (Cer et al., 2017) is a text similarity task; and QQP is a duplication detection task. There are also four natural language inference tasks in GLUE: MNLI (Williams et al., 2018); QNLI (Rajpurkar et al., 2016); RTE (Dagan et al., 2006; Bar-Haim et al., 2006; Giampiccolo et al., 2007; Bentivogli et al., 2009); and WNLI (Levesque et al., 2012). Following previous works on model distillation, we exclude STS-B and WNLI in the experiments. Dataset details are summarized in Appendix A.
Question Answering. We evaluate the proposed algorithm on two question answering datasets:
SQuAD v1.1 (Rajpurkar et al., 2016) and SQuAD v2.0 (Rajpurkar et al., 2018). These tasks are treated as a sequence labeling problem, where we predict the probability of each token being the start and end of the answer span. Dataset details can be found in Appendix A.
We compare our method with both task-agnostic and task-specific distillation methods.
In task-agnostic distillation, we pre-train a small language model through knowledge distillation, and then fine-tune on downstream tasks. The finetuning procedure also incorporates task-specific distillation for better performance.
DistilBERT (Sanh et al., 2019) pre-trains a small language model by distilling the temperature-controlled soft prediction probability.
TinyBERT (Jiao et al., 2020) is a task-agnostic distillation method that adopts layer-wise distillation.
MiniLMv1 (Wang et al., 2020) and MiniLMv2 (Wang et al., 2021) pre-train a small language model by aligning the attention distribution between the teacher model and the student model.
CoDIR (Contrastive Distillation, Sun et al. 2020a) proposes a framework that distills knowledge through intermediate Transformer layers of the teacher via a contrastive objective.
In task-specific distillation, a pre-trained language model is directly compressed and fine-tuned.
PKD (Patient Knowledge Distillation, Sun et al. 2019) proposes a method where the student patiently learns from multiple intermediate Transformer layers of the teacher.
BERT-of-Theseus (Xu et al., 2020) proposes a progressive module replacing method for knowledge distillation.
In the experiments, we use BERT-base (Devlin et al., 2019) as both the student model and the teacher model. That is, we first transform the pretrained model into a MoE model, and then apply layer-wise task-specific knowledge distillation. We set the number of experts in the MoE model to 4, and the hidden dimension of each expert is set to 768, a quarter of the hidden dimension of BERT-base. The other configurations remain unchanged. We share the top-512 important neurons among the experts (see Section 3.1). The number of effective
4https://github.com/huggingface/ transformers
RTE Acc	CoLA Mcc	MRPC F1/Acc	SST-2 Acc	QNLI Acc	QQP F1/Acc	MNLI m/mm
BERT-base	63.5	54.7	89.0/84.1	92.9	91.1	88.3/90.9	84.5/84.4
Task-agnostic
DistilBERT	59.9	51.3	87.5/-	92.7	89.2	-/88.5	82.2/-
TinyBERT (w/o aug)	72.2	42.8	88.4/-	91.6	90.5	-/90.6	83.5/-
MiniLMv1	71.5	49.2	88.4/-	92.0	91.0	-/91.0	84.0/-
MiniLMv2	72.1	52.5	88.9/-	92.4	90.8	-/91.1	84.2/-
CoDIR (pre+fine)	67.1	53.7	89.6/-	93.6	90.1	-/89.1	83.5/82.7
Task-specific
PKD	65.5	24.8	86.4/-	92.0	89.0	-/88.9	81.5/81.0
BERT-of-Theseus	68.2	51.1	89.0/-	91.5	89.5	-/89.6	82.3/-
CoDIR (fine)	65.6	53.6	89.4/-	93.6	90.4	-/89.1	83.6/82.8
Ours (task-specific)
MoEBERT	74.0	55.4	92.6/89.5	93.0	91.3	88.4/91.4	84.5/84.8