SlowGuess/ACL_2020_OCR huggingface dataset:
Zhiqing Sun $^{1*}$ , Hongkun Yu $^{2}$ , Xiaodan Song $^{2}$ , Renjie Liu $^{2}$ , Yiming Yang $^{1}$ , Denny Zhou $^{2}$
$^{1}$ Carnegie Mellon University {zhiqings, yiming}@cs.cmu.edu  $^{2}$ Google Brain {hongkuny, xiaodansong, renjieliu, dennyzhou}@google.com
Natural Language Processing (NLP) has recently achieved great success by using huge pre-trained models with hundreds of millions of parameters. However, these models suffer from heavy model sizes and high latency such that they cannot be deployed to resource-limited mobile devices. In this paper, we propose MobileBERT for compressing and accelerating the popular BERT model. Like the original BERT, MobileBERT is task-agnostic, that is, it can be generically applied to various downstream NLP tasks via simple fine-tuning. Basically, MobileBERT is a thin version of  $\mathrm{BERT}_{\mathrm{LARGE}}$ , while equipped with bottleneck structures and a carefully designed balance between self-attention and feed-forward networks. To train MobileBERT, we first train a specially designed teacher model, an inverted-bottleneck incorporated  $\mathrm{BERT}_{\mathrm{LARGE}}$  model. Then, we conduct knowledge transfer from this teacher to MobileBERT. Empirical studies show that MobileBERT is  $4.3\times$  smaller and  $5.5\times$  faster than  $\mathrm{BERT}_{\mathrm{BASE}}$  while achieving competitive results on well-known benchmarks. On the natural language inference tasks of GLUE, MobileBERT achieves a GLUE score of 77.7 (0.6 lower than  $\mathrm{BERT}_{\mathrm{BASE}}$ ), and 62 ms latency on a Pixel 4 phone. On the SQuAD v1.1/v2.0 question answering task, MobileBERT achieves a dev F1 score of 90.0/79.2 (1.5/2.1 higher than  $\mathrm{BERT}_{\mathrm{BASE}}$ ).
The NLP community has witnessed a revolution of pre-training self-supervised models. These models usually have hundreds of millions of parameters (Peters et al., 2018; Radford et al., 2018; Devlin et al., 2018; Radford et al., 2019; Yang et al., 2019). Among these models, BERT (Devlin et al., 2018)
shows substantial accuracy improvements. However, as one of the largest models ever in NLP, BERT suffers from the heavy model size and high latency, making it impractical for resource-limited mobile devices to deploy the power of BERT in mobile-based machine translation, dialogue modeling, and the like.
There have been some efforts that task-specifically distill BERT into compact models (Turc et al., 2019; Tang et al., 2019; Sun et al., 2019; Tsai et al., 2019). To the best of our knowledge, there is not yet any work for building a task-agnostic lightweight pre-trained model, that is, a model that can be generically fine-tuned on different downstream NLP tasks as the original BERT does. In this paper, we propose MobileBERT to fill this gap. In practice, task-agnostic compression of BERT is desirable. Task-specific compression needs to first fine-tune the original large BERT model into a task-specific teacher and then distill. Such a process is much more complicated (Wu et al., 2019) and costly than directly fine-tuning a task-agnostic compact model.
At first glance, it may seem straightforward to obtain a task-agnostic compact BERT. For example, one may just take a narrower or shallower version of BERT, and train it until convergence by minimizing a convex combination of the prediction loss and distillation loss (Turc et al., 2019; Sun et al., 2019). Unfortunately, empirical results show that such a straightforward approach results in significant accuracy loss (Turc et al., 2019). This may not be that surprising. It is well-known that shallow networks usually do not have enough representation power while narrow and deep networks are difficult to train.
Our MobileBERT is designed to be as deep as  $\mathrm{BERT}_{\mathrm{LARGE}}$  while each layer is made much narrower via adopting bottleneck structures and balancing between self-attention and feed-forward
This work was done when the first author was an intern at Google Brain.
Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2158-2170 July 5 - 10, 2020. ©2020 Association for Computational Linguistics
Figure 1: Illustration of three models: (a) BERT; (b) Inverted-Bottleneck BERT (IB-BERT); and (c) MobileBERT. In (b) and (c), red lines denote inter-block flows while blue lines intra-block flows. MobileBERT is trained by layer-to-layer imitating IB-BERT.
networks (Figure 1). To train MobileBERT, a deep and thin model, we first train a specially designed teacher model, an inverted-bottleneck incorporated BERTLarge model (IB-BERT). Then, we conduct knowledge transfer from IB-BERT to MobileBERT. A variety of knowledge transfer strategies are carefully investigated in our empirical studies.
Empirical evaluations show that MobileBERT is  $4.3 \times$  smaller and  $5.5 \times$  faster than BERT<sub>BASE</sub>, while it can still achieve competitive results on well-known NLP benchmarks. On the natural language inference tasks of GLUE, MobileBERT can achieve a GLUE score of 77.7, which is only 0.6 lower than BERT<sub>BASE</sub>, with a latency of 62 ms on a Pixel 4 phone. On the SQuAD v1.1/v2.0 question answering task, MobileBER obtains a dev F1 score of 90.3/80.2, which is even 1.5/2.1 higher than BERT<sub>BASE</sub>.
Recently, compression of BERT has attracted much attention. Turc et al. (2019) propose to pre-train the smaller BERT models to improve task-specific knowledge distillation. Tang et al. (2019) distill BERT into an extremely small LSTM model. Tsai et al. (2019) distill a multilingual BERT into smaller BERT models on sequence labeling tasks. Clark et al. (2019b) use several single-task BERT
models to teach a multi-task BERT. Liu et al. (2019a) distill knowledge from an ensemble of BERT models into a single BERT.
Concurrently to our work, Sun et al. (2019) distill BERT into shallower students through knowledge distillation and an additional knowledge transfer of hidden states on multiple intermediate layers. Jiao et al. (2019) propose TinyBERT, which also uses a layer-wise distillation strategy for BERT but in both pre-training and fine-tuning stages. Sanh et al. (2019) propose DistilBERT, which successfully halves the depth of BERT model by knowledge distillation in the pre-training stage and an optional fine-tuning stage.
In contrast to these existing literature, we only use knowledge transfer in the pre-training stage and do not require a fine-tuned teacher or data augmentation (Wu et al., 2019) in the down-stream tasks. Another key difference is that these previous work try to compress BERT by reducing its depth, while we focus on compressing BERT by reducing its width, which has been shown to be more effective (Turc et al., 2019).
3 MobileBERT
In this section, we present the detailed architecture design of MobileBERT and training strategies to efficiently train MobileBERT. The specific model settings are summarized in Table 1. These settings are obtained by extensive architecture search experiments which will be presented in Section 4.1.
1The code and pre-trained models will be available at https://github.com/google-research/google-research/tree/master/mobilebert.
BERTLARGE	BERTBASE	IB-BERTLARGE	MobileBERT	MobileBERTTINY
embedding	1024	768	128
no-op	no-op	3-convolution
1024	768	512
body	Linear	hinput	[ (1024)	[ (768)	[ (512)	[ (512)
houtput	16	12	1024	128
MHA	hinput	[ 1024)	[ 768)	[ 512)	[ 128)
#Head	16	12	4	128
houtput	1024	768)	1024	128
FFN	hinput	[ 4096)	[ 768)	[ 4096)	[ 128)
hFN	1024	768)	1024	128
houtput	1024	768)	1024	128
Linear	hinput
houtput
#Params	334M	109M	293M	25.3M	15.1M
Table 1: The detailed model settings of a few models.  $\mathrm{h}_{\text{inter}}$ ,  $\mathrm{h}_{\text{FFN}}$ ,  $\mathrm{h}_{\text{embedding}}$ , #Head and #Params denote the inter-block hidden size (feature map size), FFN intermediate size, embedding table size, the number of heads in multi-head attention, and the number of parameters, respectively.
3.1 Bottleneck and Inverted-Bottleneck
The architecture of MobileBERT is illustrated in Figure 1(c). It is as deep as  $\mathrm{BERT}_{\mathrm{LARGE}}$ , but each building block is made much smaller. As shown in Table 1, the hidden dimension of each building block is only 128. On the other hand, we introduce two linear transformations for each building block to adjust its input and output dimensions to 512. Following the terminology in (He et al., 2016), we refer to such an architecture as bottleneck.
It is challenging to train such a deep and thin network. To overcome the training issue, we first construct a teacher network and train it until convergence, and then conduct knowledge transfer from this teacher network to MobileBERT. We find that this is much better than directly training MobileBERT from scratch. Various training strategies will be discussed in a later section. Here, we introduce the architecture design of the teacher network which is illustrated in Figure 1(b). In fact, the teacher network is just  $\mathrm{BERT}_{\mathrm{LARGE}}$  while augmented with inverted-bottleneck structures (Sandler et al., 2018) to adjust its feature map size to 512. In what follows, we refer to the teacher network as IB-BERT $_{\mathrm{LARGE}}$ . Note that IB-BERT and MobileBERT have the same feature map size which is 512. Thus, we can directly compare the layer-wise output difference between IB-BERT and MobileBERT. Such a direct comparison is needed in our knowledge transfer strategy.
It is worth pointing out that the simultaneously introduced bottleneck and inverted-bottleneck structures result in a fairly flexible architecture design. One may either only use the bottlenecks for MobileBERT (correspondingly the
teacher becomes  $\mathrm{BERT}_{\mathrm{LARGE}}$  or only the inverted bottlenecks for IB-BERT (then there is no bottleneck in MobileBERT) to align their feature maps. However, when using both of them, we can allow IB-BERT $_{\mathrm{LARGE}}$  to preserve the performance of  $\mathrm{BERT}_{\mathrm{LARGE}}$  while having MobileBERT sufficiently compact.
3.2 Stacked Feed-Forward Networks
A problem introduced by the bottleneck structure of MobileBERT is that the balance between the Multi-Head Attention (MHA) module and the Feed-Forward Network (FFN) module is broken. MHA and FFN play different roles in the Transformer architecture: The former allows the model to jointly attend to information from different subspaces, while the latter increases the non-linearity of the model. In original BERT, the ratio of the parameter numbers in MHA and FFN is always 1:2. But in the bottleneck structure, the inputs to the MHA are from wider feature maps (of inter-block size), while the inputs to the FFN are from narrower bottlenecks (of intra-block size). This results in that the MHA modules in MobileBERT relatively contain more parameters.
To fix this issue, we propose to use stacked feedforward networks in MobileBERT to re-balance the relative size between MHA and FFN. As illustrated in Figure 1(c), each MobileBERT layer contains one MHA but several stacked FFN. In MobileBERT, we use 4 stacked FFN after each MHA.
3.3 Operational Optimizations
By model latency analysis $^2$ , we find that layer normalization (Ba et al., 2016) and gelu activation (Hendrycks and Gimpel, 2016) accounted for a considerable proportion of total latency. Therefore, we propose to replace them with new operations in our MobileBERT.
Remove layer normalization We replace the layer normalization of a  $n$ -channel hidden state  $\mathbf{h}$  with an element-wise linear transformation:
$$
\operatorname {N o N o r m} (\mathbf {h}) = \gamma \circ \mathbf {h} + \beta , \tag {1}
$$
where  $\gamma, \beta \in \mathbb{R}^n$  and  $\circ$  denotes the Hadamard product. Please note that NoNorm has different properties from LayerNorm even in test mode since the original layer normalization is not a linear operation for a batch of vectors.
Use relu activation We replace the gelu activation with simpler relu activation (Nair and Hinton, 2010).
3.4 Embedding Factorization
The embedding table in BERT models accounts for a substantial proportion of model size. To compress the embedding layer, as shown in Table 1, we reduce the embedding dimension to 128 in MobileBERT. Then, we apply a 1D convolution with kernel size 3 on the raw token embedding to produce a 512 dimensional output.
3.5 Training Objectives
We propose to use the following two knowledge transfer objectives, i.e., feature map transfer and attention transfer, to train MobileBERT. Figure 1 illustrates the proposed layer-wise knowledge transfer objectives. Our final layer-wise knowledge transfer loss  $\mathcal{L}_{KT}^{\ell}$  for the  $\ell^{th}$  layer is a linear combination of the two objectives stated below:
Feature Map Transfer (FMT) Since each layer in BERT merely takes the output of the previous layer as input, the most important thing in layer-wise knowledge transfer is that the feature maps of each layer should be as close as possible to those of the teacher. In particular, the mean squared error between the feature maps of the MobileBERT
student and the IB-BERT teacher is used as the knowledge transfer objective:
$$
\mathcal {L} _ {F M T} ^ {\ell} = \frac {1}{T N} \sum_ {t = 1} ^ {T} \sum_ {n = 1} ^ {N} (H _ {t, \ell , n} ^ {t r} - H _ {t, \ell , n} ^ {s t}) ^ {2}, \quad (2)
$$
where  $\ell$  is the index of layers,  $T$  is the sequence length, and  $N$  is the feature map size. In practice, we find that decomposing this loss term into normalized feature map discrepancy and feature map statistics discrepancy can help stabilize training.
Attention Transfer (AT) The attention mechanism greatly boosts the performance of NLP and becomes a crucial building block in Transformer and BERT (Clark et al., 2019a; Jawahar et al., 2019). This motivates us to use self-attention maps from the well-optimized teacher to help the training of MobileBERT in augmentation to the feature map transfer. In particular, we minimize the KL-divergence between the per-head self-attention distributions of the MobileBERT student and the IB-BERT teacher:
$$
\mathcal {L} _ {A T} ^ {\ell} = \frac {1}{T A} \sum_ {t = 1} ^ {T} \sum_ {a = 1} ^ {A} D _ {K L} \left(a _ {t, \ell , a} ^ {t r} | | a _ {t, \ell , a} ^ {s t}\right), \tag {3}
$$
where  $A$  is the number of attention heads.
Pre-training Distillation (PD) Besides layer-wise knowledge transfer, we can also use a knowledge distillation loss when pre-training MobileBERT. We use a linear combination of the original masked language modeling (MLM) loss, next sentence prediction (NSP) loss, and the new MLM Knowledge Distillation (KD) loss as our pretraining distillation loss:
$$
\mathcal {L} _ {P D} = \alpha \mathcal {L} _ {M L M} + (1 - \alpha) \mathcal {L} _ {K D} + \mathcal {L} _ {N S P}, \tag {4}
$$
where  $\alpha$  is a hyperparameter in  $(0,1)$ .
3.6 Training Strategies
Given the objectives defined above, there can be various combination strategies in training. We discuss three strategies in this paper.
Auxiliary Knowledge Transfer In this strategy, we regard intermediate knowledge transfer as an auxiliary task for knowledge distillation. We use a single loss, which is a linear combination of knowledge transfer losses from all layers as well as the pre-training distillation loss.
2A detailed analysis of effectiveness of operational optimizations on real-world inference latency can be found in Section 4.6.1.
Figure 2: Diagrams of (a) auxiliary knowledge transfer (AKT), (b) joint knowledge transfer (JKT), and (c) progressive knowledge transfer (PKT). Lighter colored blocks represent that they are frozen in that stage.
Joint Knowledge Transfer However, the intermediate knowledge of the IB-BERT teacher (i.e. attention maps and feature maps) may not be an optimal solution for the MobileBERT student. Therefore, we propose to separate these two loss terms, where we first train MobileBERT with all layer-wise knowledge transfer losses jointly, and then further train it by pre-training distillation.
Progressive Knowledge Transfer One may also concern that if MobileBERT cannot perfectly mimic the IB-BERT teacher, the errors from the lower layers may affect the knowledge transfer in the higher layers. Therefore, we propose to progressively train each layer in the knowledge transfer. The progressive knowledge transfer is divided into  $L$  stages, where  $L$  is the number of layers.
Diagram of three strategies Figure 2 illustrates the diagram of the three strategies. For joint knowledge transfer and progressive knowledge transfer, there is no knowledge transfer for the beginning embedding layer and the final classifier in the layerwise knowledge transfer stage. They are copied from the IB-BERT teacher to the MobileBERT student. Moreover, for progressive knowledge transfer, when we train the  $\ell^{th}$  layer, we freeze all the trainable parameters in the layers below. In practice, we can soften the training process as follows. When training a layer, we further tune the lower layers with a small learning rate rather than entirely freezing them.
In this section, we first present our architecture search experiments which lead to the model settings in Table 1, and then present the empirical
#Params	hinter	hintra	#Head	SQuAD
(a)	356M	1024	1024	16	88.2
(b)	325M	768	1024	16	88.6
(c)	293M	512	1024	16	88.1
(d)	276M	384	1024	16	87.6
(e)	262M	256	1024	16	87.0
(f)	293M	512	1024	4	88.3
(g)	92M	512	512	4	85.8
(h)	33M	512	256	4	84.8
(i)	15M	512	128	4	82.0
Table 2: Experimental results on SQuAD v1.1 dev F1 score in search of good model settings for the IB-BERTLARGE teacher. The number of layers is set to 24 for all models.
results on benchmarks from MobileBERT and various baselines.
4.1 Model Settings
We conduct extensive experiments to search good model settings for the IB-BERT teacher and the MobileBERT student. We start with SQuAD v1.1 dev F1 score as the performance metric in the search of model settings. In this section, we only train each model for 125k steps with 2048 batch size, which halves the training schedule of original BERT (Devlin et al., 2018; You et al., 2019).
Architecture Search for IB-BERT Our design philosophy for the teacher model is to use as small inter-block hidden size (feature map size) as possible, as long as there is no accuracy loss. Under this guideline, we design experiments to manipulate the inter-block size of a BERT $_{\text{LARGE}}$ -sized IB-BERT, and the results are shown in Table 2 with labels (a)-(e). We can see that reducing the interblock hidden size doesn't damage the performance
h_intra	#Head (#Params)	#FFN (#Params)	SQuAD
192	6	(8M)	1	(7M)	82.6
160	5	(6.5M)	2	(10M)	83.4
128	4	(5M)	4	(12.5M)	83.4
96	3	(4M)	8	(14M)	81.6
Table 3: Experimental results on SQuAD v1.1 dev F1 score in search of good model settings for the MobileBERT student. The number of layers is set to 24 and the inter-block hidden size is set to 512 for all models.
of BERT until it is smaller than 512. Hence, we choose IB-BERT $_{\text{LARGE}}$  with its inter-block hidden size being 512 as the teacher model.
One may wonder whether we can also shrink the intra-block hidden size of the teacher. We conduct experiments and the results are shown in Table 2 with labels (f)-(i). We can see that when the intra-block hidden size is reduced, the model performance is dramatically worse. This means that the intra-block hidden size, which represents the representation power of non-linear modules, plays a crucial role in BERT. Therefore, unlike the interblock hidden size, we do not shrink the intra-block hidden size of our teacher model.
Architecture Search for MobileBERT We seek a compression ratio of  $4 \times$  for  $\mathrm{BERT}_{\mathrm{BASE}}$ , so we design a set of MobileBERT models all with approximately 25M parameters but different ratios of the parameter numbers in MHA and FFN to select a good MobileBERT student model. Table 3 shows our experimental results. They have different balances between MHA and FFN. From the table, we can see that the model performance reaches the peak when the ratio of parameters in MHA and FFN is  $0.4 \sim 0.6$ . This may justify why the original Transformer chooses the parameter ratio of MHA and FFN to 0.5.
We choose the architecture with 128 intra-block hidden size and 4 stacked FFNs as the MobileBERT student model in consideration of model accuracy and training efficiency. We also accordingly set the number of attention heads in the teacher model to 4 in preparation for the layer-wise knowledge transfer. Table 1 demonstrates the model settings of our IB-BERT $_{\text{LARGE}}$  teacher and MobileBERT student.
One may wonder whether reducing the number of heads will harm the performance of the teacher model. By comparing (a) and (f) in Table 2, we can see that reducing the number of heads from 16 to 4
does not affect the performance of IB-BERTLARGE.
Following BERT (Devlin et al., 2018), we use the BooksCorpus (Zhu et al., 2015) and English Wikipedia as our pre-training data. To make the IB-BERTLarge teacher reach the same accuracy as original BERTLarge, we train IB-BERTLarge on 256 TPU v3 chips for 500k steps with a batch size of 4096 and LAMB optimizer (You et al., 2019). For a fair comparison with the original BERT, we do not use training tricks in other BERT variants (Liu et al., 2019b; Joshi et al., 2019). For MobileBERT, we use the same training schedule in the pre-training distillation stage. Additionally, we use progressive knowledge transfer to train MobileBERT, which takes additional 240k steps over 24 layers. In ablation studies, we halve the pretraining distillation schedule of MobileBERT to accelerate experiments. Moreover, in the ablation study of knowledge transfer strategies, for a fair comparison, joint knowledge transfer and auxiliary knowledge transfer also take additional 240k steps.
For the downstream tasks, all reported results are obtained by simply fine-tuning MobileBERT just like what the original BERT does. To fine-tune the pre-trained models, we search the optimization hyperparameters in a search space including different batch sizes (16/32/48), learning rates ((1-10) * e-5), and the number of epochs (2-10). The search space is different from the original BERT because we find that MobileBERT usually needs a larger learning rate and more training epochs in fine-tuning. We select the model for testing according to their performance on the development (dev) set.
4.3 Results on GLUE
The General Language Understanding Evaluation (GLUE) benchmark (Wang et al., 2018) is a collection of 9 natural language understanding tasks. We compare MobileBERT with  $\mathrm{BERT}_{\mathrm{BASE}}$  and a few state-of-the-art pre-BERT models on the GLUE leaderboard: OpenAI GPT (Radford et al., 2018) and ELMo (Peters et al., 2018). We also compare with three recently proposed compressed BERT models: BERT-PKD (Sun et al., 2019), and DistilBERT (Sanh et al., 2019). To further show the advantage of MobileBERT over recent small BERT models, we also evaluate a smaller variant of our
<sup>3</sup>https://gluebenchmark.com/leaderboard
#Params #FLOPS Latency	CoLA SST-2 MRPC STS-B QQP MNLI-m/mm QNLI RTE
8.5k 67k 3.7k 5.7k 364k 393k 108k 2.5k
ELMo-BiLSTM-Attn	-	-	-	33.6	90.4	84.4	72.3	63.1	74.1/74.5	79.8	58.9	70.0
OpenAI GPT	109M	-	-	47.2	93.1	87.7	84.8	70.1	80.7/80.6	87.2	69.1	76.9
BERTBASE	109M	22.5B	342 ms	52.1	93.5	88.9	85.8	71.2	84.6/83.4	90.5	66.4	78.3
BERTBASE-6L-PKD*	66.5M	11.3B	-	-	92.0	85.0	-	70.7	81.5/81.0	89.0	65.5	-
BERTBASE-4L-PKD†*	52.2M	7.6B	-	24.8	89.4	82.6	79.8	70.2	79.9/79.3	85.1	62.3	-
BERTBASE-3L-PKD*	45.3M	5.7B	-	-	87.5	80.7	-	68.1	76.7/76.3	84.7	58.2	-
DistilBERTBASE-6L†	62.2M	11.3B	-	-	92.0	85.0	70.7	81.5/81.0	89.0	65.5	-
DistilBERTBASE-4L†	52.2M	7.6B	-	32.8	91.4	82.4	76.1	68.5	78.9/78.0	85.2	54.1	-
TinyBERT*	14.5M	1.2B	-	43.3	92.6	86.4	79.9	71.3	82.5/81.8	87.7	62.9	75.4
MobileBERTTINY	15.1M	3.1B	40 ms	46.7	91.7	87.9	80.1	68.9	81.5/81.6	89.5	65.1	75.8
MobileBERT	25.3M	5.7B	62 ms	50.5	92.8	88.8	84.4	70.2	83.3/82.6	90.6	66.2	77.7
MobileBERT w/o OPT	25.3M	5.7B	192 ms	51.1	92.6	88.8	84.8	70.5	84.3/83.4	91.6	70.4	78.5