SlowGuess/ACL_2020_OCR huggingface dataset:
Ashish Khetan
khetan@amazon.com
Zohar Karnin
zkarnin@amazon.com
Transformers (Vaswani et al., 2017) have gradually become a key component for many state-of-the-art natural language representation models. A recent Transformer based model-BERT (Devlin et al., 2018) achieved state-of-the-art results on various natural language processing tasks, including GLUE, SQuAD v1.1, and SQuAD v2.0. This model however is computationally prohibitive and has a huge number of parameters. In this work we revisit the architecture choices of BERT in efforts to obtain a lighter model. We focus on reducing the number of parameters yet our methods can be applied towards other objectives such FLOPs or latency. We show that much efficient light BERT models can be obtained by reducing algorithmically chosen correct architecture design dimensions rather than reducing the number of Transformer encoder layers. In particular, our schuBERT gives  $6.6\%$  higher average accuracy on GLUE and SQuAD datasets as compared to BERT with three encoder layers while having the same number of parameters.
Transformer (Vaswani et al., 2017) based models have achieved state-of-the-art performance for many natural language processing tasks (Dai and Le, 2015; Peters et al., 2018; Radford et al., 2018; Howard and Ruder, 2018). These include machine translation (Vaswani et al., 2017; Ott et al., 2018), question-answering tasks (Devlin et al., 2018), natural language inference (Bowman et al., 2015; Williams et al., 2017) and semantic role labeling (Strubell et al., 2018).
A recent Transformer based model BERT (Devlin et al., 2018) achieved state-of-the-art results on various natural language processing tasks including GLUE, SQuAD v1.1 and SQuAD v2.0. BERT's model architecture is a multi-layer bidirectional
Transformer encoder based on the original implementation described in Vaswani et al. (2017).
Following the seminal results obtained by the BERT model, several follow up studies explored methods for improving them further. XLNet (Yang et al., 2019) adds autoregressive capabilities to BERT, improving its quality, though at the cost of additional compute requirements. RoBERTa (Liu et al., 2019) modifies the training procedure of BERT and provides pre-training methods that significantly improve its performance. Two notable papers exploring the architecture design of the BERT are following. Michel et al. (2019) examines the importance of attention heads in BERT architecture, highlighting scenarios where attention heads may be pruned. The main objective of the paper is to provide techniques for pruning attention head, and as such the amount of experiments performed on BERT is limited to a single task (MNLI). ALBERT (Lan et al., 2019) proposes two methods for reducing the number of parameters in BERT. The first is via parameter sharing across layers, and the second is by factorizing the embedding layers. We note (this was mentioned in the conclusion section of the paper) that while these methods are efficient in reducing the number of parameters used by the model, they do not help in reducing its latency.
These studies provide some advancement towards a more efficient architecture design for BERT but leave much to be explored. In this paper we take a broader approach examining multiple design choices. We parameterize each layer of BERT by five different dimensions, as opposed to Devlin et al. (2018) that parameterizes a layer with two dimensions and suggests a fixed value for the remaining three. We then (pre-)train multiple variants of BERT with different values chosen for these dimensions by applying pruning-based architecture search technique that jointly optimizes the architecture of the model with the objective of minimizing
Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2807-2818 July 5 - 10, 2020. ©2020 Association for Computational Linguistics
both the pre-training loss and the number of model parameters. Our experiments result in the following findings:
- The ratio of the architecture design dimensions within a BERT encoder layer can be modified to obtain a layer with better performance. Transformer design dimensions suggested in Vaswani et al. (2017) are suboptimal.
- When we aim to obtain a computationally lighter model, using a 'tall and narrow' architecture provides better performance than a 'wide and shallow' architecture.
- The fully-connected component applied to each token separately plays a much more significant role in the top layers as compared to the bottom layers.
Following BERT's notations, we use  $\ell$  to denote the number of encoder layers (i.e. Transformer blocks),  $h$  to denote the hidden size, and  $a$  to denote the number of self attention heads. The BERT paper (Devlin et al., 2018) primarily reports results on two models:  $\mathrm{BERT}_{\mathrm{BASE}}$  ( $\ell = 12, h = 768, a = 12$ ) and  $\mathrm{BERT}_{\mathrm{LARGE}}$  ( $\ell = 24, h = 1024, a = 16$ ). BERT base has 108M parameters and BERT large has 340M parameters. Though BERT large achieves higher accuracy than BERT base, due to its prohibitively large size it finds limited use in practice. Since BERT base achieves higher accuracy compared to previous state-of-the-art models- Pre-OpenAI SOTA, BiLSTM+ELMo+Attn and OpenAI GPT-on most of the benchmark datasets, it is widely used in practice. BERT base and OpenAI GPT have the same number of model parameters.
Given its broad adoption for NLP tasks, an immediate question is: can we reduce the size of BERT base without incurring any significant loss in accuracy? The BERT paper (Devlin et al., 2018) provides an ablation study, Table 1, over the number of model parameters by varying the number of layers  $\ell$ , the hidden size  $h$ , and the number of attention heads  $a$ . It can be observed that the accuracy decreases drastically when the number of encoder layers  $\ell$  is reduced, and also when the number of attention heads is reduced. We ask the following question: are there any other design dimensions that can be reduced without incurs huge loss in accuracy?
Design dimensions	Dev Set Accuracy
#l	#a	#M	MNLI	MRPC	SST-2
3	12	45	77.9	79.8	88.4
6	3	55	80.6	82.2	90.7
6	12	66	81.9	84.8	91.3
BERT base
12	12	108	84.4	86.7	92.9
Table 1: Ablation study over BERT model size, Table 6 in Devlin et al. (2018). #M denotes number of model parameters in millions. hidden size,  $h = {768}$  .
As noted above, the three primary design dimensions of the BERT architecture are the number of encoder layers  $\ell$ , the hidden size  $h$ , and the number of attention heads  $a$ . BERT's Transformer encoder layers are based on the original Transformer implementation described in Vaswani et al. (2017). Vaswani et al. (2017) fixed dimension of key, query, and value in multi-head attention, and filter dimension in feed-forward networks as a function of the hidden size and the number of attention heads. However, these are variable design dimensions and can be optimized. Moreover, BERT architecture uses the same number of attention heads for all the encoder layers and hence all the layers are identical. In this work, we jointly optimize all these design dimensions of BERT architecture while allowing each encoder layer to have different design dimensions.
In order to explore the parameter space efficiently we chose to optimize the design dimensions in a pruning framework rather than launching a pretraining job for each of these choices. This allows a speedup of several orders of magnitude that is crucial in order to obtain meaningful conclusions. We parameterize the different dimensions one can modify and jointly optimize them with a mixed target of both accuracy and parameter reduction. We look at how the accuracy of BERT evolves on various downstream datasets like GLUE, SQuAD v1.1, and SQuAD v2.0 when we reduce the model size via an optimization procedure.
There is a vast literature on pruning trained neural networks. Starting with the classical works LeCun et al. (1990); Hassibi and Stork (1993) in the early 90's to the recent works Han et al. (2015), pruning deep neural networks has received a lot of attention. There have been two orthogonal ap
proaches in pruning networks: structured pruning (Li et al., 2016; Molchanov et al., 2016) and unstructured pruning (Anwar et al., 2017). Structured pruning gives smaller architecture whereas unstructured pruning gives sparse model parameters. In natural language processing, Murray and Chiang (2015) explored structured pruning in feed-forward language models. See et al. (2016) and Kim and Rush (2016) provided pruning approaches for machine translation. A closely related line of work is Neural Architecture Search (NAS). It aims to efficiently search the space of architectures (Pham et al., 2018; Liu et al., 2018; Singh et al., 2019). Quantization is another technique to reduce the model size. This is done by quantizing the model parameters to binary (Rastegari et al., 2016; Hubara et al., 2017), ternary (Zhu et al., 2016), or 4 or 8 bits per parameter (Han et al., 2015).
Recently published DistilBERT (Sanh et al., 2019) shows that a BERT model with fewer number of layers can be efficiently pre-trained using knowledge distillation to give much higher accuracy as compared to the same model pre-trained in a regular way. We note that the distillation technique is complimentary to our work and our schuBERTs can be pre-trained using distillation to boost their accuracy. The ablation study in Table 1, BERT (Devlin et al., 2018), and the above explained works (Michel et al., 2019; Lan et al., 2019) look at the problem of reducing the BERT model size by reducing one or the other design dimensions - number of encoder layers, hidden size, number of attention heads, and embedding size - in isolation and in a sub-optimal way. In this work, we address this problem comprehensively.
4 The Elements of BERT
In this section, we present detailed architecture of the original BERT model and explain which design dimensions of it can be optimized. Figure 1 shows BERT pre-training architecture. First, the tokenized inputs are embedded into a vector of dimension  $h$  through an embedding layer  $E$ . The embedded inputs pass through a sequence of encoder layers 1 to  $\ell$ . Each encoder layer is identical in its architecture. The output of the last encoder layer is decoded using the same embedding layer  $E$  and softmax cross-entropy loss is computed on the masked tokens. A special token CLS from the last encoder layer is used to compute next-sentence-prediction (NSP) loss. For further details of the loss
Figure 1: BERT pre-training
BERT-base
number of encoder layers	l	12
hidden size	h	768
number of self-attention heads	a	12
feed forward dimension	f	4h
key-query dimension for attention	k	h/a
value dimension for attention	v	h/a
Table 2: Elements of BERT
corresponding to masked tokens and the NSP loss, we refer the readers to the BERT paper (Devlin et al., 2018).
We follow BERT notation conventions and denote the number of encoder layers as  $\ell$ , the hidden size as  $h$ , and the number of attention heads as  $a$ . Following the original Transformer implementation described in Vaswani et al. (2017) BERT sets key-query dimension for multi-head attention  $k$  to  $h / a$ . Following the same Transformer implementation it sets value dimension for multi-head attention  $v$  equal to  $k$ , and feed-forward filter size  $f$  equal to  $4h$ . In total, there are three design dimensions in BERT- $\ell$ ,  $h$  and  $a$ , they are listed in Table 2. For BERT base, the number of encoder layers  $\ell$  is set to 12, the hidden size  $h$  is set to 768, and the number of attention heads  $a$  is set to 12. The other three dimensions  $f, k, v$  are function of  $h$  and  $a$ . Further, each encoder layer of BERT is identical and uses same value of  $a, f, k, v$ .
First of all, BERT has no architectural constraint that requires all the encoder layers to be identical. This aspect of design can be optimized and in full-generality it might result in highly non-identical layers. This implies that a generalized
schuBERT
l	l
h	h
a	a1,a2,···,al
f	f1,f2,···,fl
k	k1,k2,···,kl
v	v1,v2,···,vl
Table 3: Elements of schuBERT
BERT will have  $a_1, a_2, \dots, a_\ell$  number of heads,  $f_1, f_2, \dots, f_\ell$  filter sizes in the feed forward networks,  $k_1, k_2, \dots, k_\ell$  key sizes and  $v_1, v_2, \dots, v_\ell$  value sizes in the attention heads, in the layers  $1, 2, \dots, \ell$  respectively. Table 3 lists all the design dimensions of BERT that can be optimized without changing the architecture. Note that we abuse the term architecture to refer to the entire BERT network and the layer operations except sizes of the parameter matrices. In this work, our goal is to optimize (by pruning) all these dimensions to maximize accuracy for a given size of the model. We refer the BERT with optimized dimensions as schuBERT-Size Constricted Hidden Unit BERT.
Now, we show which parameter matrices are tied with each of these design dimensions. Each design dimension is tied with more than one parameter matrix. This is explained by providing a detail view of an encoder cell of the BERT.
Figure 2 shows architecture of an encoder layer of BERT. The notations in the figure have subscript 1 that represent first encoder layer. Input to an encoder layer is the hidden representation of a token which is of dimension  $h$ . Input first goes through a multi-head attention cell. Note that multi-head attention cell processes hidden representation of all the tokens in a combined way. For simplicity, in Figure 2 we have shown only one hidden representation.
The multi-head attention cell consists of three parameter tensors, namely - key  $K_{1}$ , query  $Q_{1}$  and value  $V_{1}$ .  $K_{1}$  is of size  $k_{1} \times a_{1} \times h$ . Key vector for each head of the attention is of dimension  $k_{1}$  and  $a_{1}$  represents the number of heads. Hidden representation of dimension  $h$  is projected on the key tensor  $K_{1}$  to get  $a_{1}$  key vectors each of dimension  $k_{1}$ . Similarly the query tensor  $Q_{1}$  is used to get  $a_{1}$  query vectors each of dimension  $k_{1}$  for  $a_{1}$  heads of the multi-head attention cell. The value tensor  $V_{1}$  is of dimension  $v_{1} \times a_{1} \times h$ . The hidden representation is projected on the value tensor  $V_{1}$  to get  $a_{1}$
value vectors each of dimension  $v_{1}$ . Note that  $k_{1}$  and  $v_{1}$  can be different. The inner product of key and query vectors after passing through softmax layer give weights for combining value vectors. For details of multi-head attention cell we refer the readers to Vaswani et al. (2017). In nutshell, using three parameter tensors-  $K_{1}, Q_{1}, V_{1}$ , a multi-head attention cell transforms hidden representation of size  $h$  to a vector of dimension  $(v_{1} \times a_{1})$ . This vector is projected back to the same dimension  $h$  through a proj matrix  $P_{1}$ . Which is then added element-wise to the hidden representation that was input to the encoder cell and layer norm is applied on the addition. The output is passed sequentially through two fully-connected layers namely  $D_{1}$  and  $G_{1}$ .  $D_{1}$  consists of a parameter matrix of dimension  $f_{1} \times h$  and  $G_{1}$  consists of a parameter matrix of dimension  $h \times f_{1}$ . The output of  $G_{1}$  is added element-wise to the input of  $D_{1}$  and layer norm is applied to it. This is the output of the encoder cell and is input to the next encoder cell.
Figure 2: An encoder layer of schuBERT
The color coding in Figure 2 shows which vectors need to be of the same dimension. The hidden representation size  $h$  needs to be same throughout all the encoder layers. In a multi-head attention cell, in each head key and query vectors must have the same dimension. Therefore, key and query tensors,  $K_{1}, Q_{1}$  must be of the same size  $k_{1} \times a_{1} \times h$ . The value vector can be of different dimension  $v_{1}$ . Therefore the value tensor  $V_{1}$  should be of dimension  $v_{1} \times a_{1} \times h$ . Further, the filter size  $f_{1}$  in the two fully-connected layers  $D_{1}, G_{1}$  is a variable and can take any integral value.
Keeping aligned with the BERT and the subse
quent improvements such as XLNet (Yang et al., 2019) and RoBERTa (Liu et al., 2019), we set the WordPiece embedding size  $e$  equal to the hidden layer size  $h$ , i.e.  $e \equiv h$ . However, factorization of the embedding matrix can be incorporated as demonstrated in ALBERT (Lan et al., 2019).
5 Optimization Method
We optimize BERT design dimensions listed in Table 3 by pruning the original BERT base architecture. All the design dimensions are upper bounded by their original value in the BERT base as given in the Table 2. Since we keep the architecture same, that is we do not remove any layer, the design dimensions are lower bounded by one.
For each design dimension that we seek to optimize, we introduce a prune-parameter vector  $\alpha$  of size equal to the original dimension. We take pretrained original BERT base network, and multiply all the parameter tensors/matrices that are associated with the particular design dimension with the corresponding prune-parameter vector. For example, filter size of the feed-forward layer in the first encoder layer is  $f_{1} = 3072$ . To optimize  $f_{1}$ , we introduce a prune-parameter vector  $\alpha_{f_1} \in \mathbf{R}^{3072}$  and initialize it with all ones. In the original BERT base, the two parameter matrices  $D_{1}$  and  $G_{1}$  are associated with the design dimension  $f_{1}$ . We replace  $D_{1}$  by  $\mathrm{diag}(\alpha_{f_1}) \cdot D_1$  and  $G_{1}$  by  $G_{1} \cdot \mathrm{diag}(\alpha_{f_{1}})$  in the BERT pre-trained model.
Table 4 lists all the prune parameters. Table 5 lists all the parameter tensors/matrices for which design dimensions are optimized by multiplying prunable parameters on all the sides. key and query tensors  $K_{i}, Q_{i}$  for  $i \in \{1,2,\dots,\ell\}$  are multiplied on all the three sides with prunable parameters corresponding to key-vector, number of attention heads, and hidden size. Similarly multiplications are performed on value tensor  $V_{i}$  with a different value-vector prunable parameter. proj tensor has same multiplication as value tensor. The two feedforward matrices  $D_{i}, G_{i}$  have same multiplications. We denote the so obtained prunable tensors with tilde on their top. Note that we do not have prune parameters for pruning encoder layers. We find the optimal number of encoder layers  $\ell$  by running experiments for different values of  $\ell$ .
Our approach is to optimally find which individual elements of prunable parameters  $\{\alpha_h, \{\alpha_{a_i}, \alpha_{v_i}, \alpha_{k_i}, \alpha_{f_i}\}_{i \in [\ell]}\}$  can be set to zero while incurring minimal increase in the pre
$$
\alpha_ {h} \in \mathbf {R} ^ {h}
$$
$$
\{\alpha_ {f _ {i}} \in \mathbf {R} ^ {f} \} _ {i = 1, 2, \dots , \ell}
$$
$$
\{\alpha_ {a _ {i}} \in \mathbf {R} ^ {a} \} _ {i = 1, 2, \dots , \ell}
$$
$$
\{\alpha_ {k _ {i}} \in \mathbf {R} ^ {k} \} _ {i = 1, 2, \dots , \ell}
$$
$$
\{\alpha_ {v _ {i}} \in \mathbf {R} ^ {v} \} _ {i = 1, 2, \dots , \ell}
$$
Table 4: Prunable parameters.
$$
K _ {i} \rightarrow K _ {i} [ \operatorname {d i a g} \left(\alpha_ {k _ {i}}\right) \operatorname {d i a g} \left(\alpha_ {a _ {i}}\right) \operatorname {d i a g} \left(\alpha_ {h}\right) ] \equiv \widetilde {K _ {i}}
$$
$$
Q _ {i} \rightarrow Q _ {i} \left[ \operatorname {d i a g} \left(\alpha_ {k _ {i}}\right) \operatorname {d i a g} \left(\alpha_ {a _ {i}}\right) \operatorname {d i a g} \left(\alpha_ {h}\right)\right] \equiv Q _ {i}
$$
$$
V _ {i} \rightarrow V _ {i} \left[ \operatorname {d i a g} \left(\alpha_ {v _ {i}}\right) \operatorname {d i a g} \left(\alpha_ {a _ {i}}\right) \operatorname {d i a g} \left(\alpha_ {h}\right)\right] \equiv \widetilde {V _ {i}}
$$
$$
P _ {i} \rightarrow P _ {i} [ \operatorname {d i a g} (\alpha_ {h}) \operatorname {d i a g} (\alpha_ {v _ {i}}) \operatorname {d i a g} (\alpha_ {a _ {i}}) ] \equiv \tilde {P} _ {i}
$$
$$
D _ {i} \rightarrow D _ {i} \left[ \operatorname {d i a g} \left(\alpha_ {f _ {i}}\right) \operatorname {d i a g} \left(\alpha_ {h}\right)\right] \equiv \bar {D} _ {i}
$$
$$
G _ {i} \rightarrow G _ {i} [ \operatorname {d i a g} (\alpha_ {h}) \operatorname {d i a g} (\alpha_ {f _ {i}}) ] \equiv G _ {i}
$$
Table 5: Prunable BERT parameter matrices/tensors.
training loss. After we have sparse prunable parameter vectors, we remove the corresponding rows/columns from the BERT parameter matrices  $\{K_i,Q_i,V_i,P_i,D_i,G_i\}_{i\in [\ell ]}$  , and get a smaller/faster BERT model. Below we explain the algorithm to find the sparse prunable parameters.
We start with the pre-trained BERT base trained on BooksCorpus (800M words) and English Wikipedia (2500M words) following the BERT pretraining procedure given in Devlin et al. (2018). Particularly, we minimize the loss given in Equation (1) to learn the optimal parameter tensors  $\{K_i, Q_i, V_i, P_i, D_i, G_i\}_{i \in [\ell]}$  and the embedding matrix  $E$ . Next, we introduce the prunable parameters given in Table 4 and initialize them with all ones. We create prunable BERT parameter matrices by multiplying the prunable parameters to the learned BERT parameter matrices, as given in Table 5. Then, we optimize the prunable parameters  $\alpha$ ’s while fixing the learned parameters matrices as given in Equation 2. In addition to the MLM and NSP loss, we add sparsity inducing loss on the prunable parameters with a regularization coefficient  $\gamma$ . It is well known that  $\ell_1$  penalty induces sparsity in the parameters. Further, since our goal is to minimize the number of parameters, to account for the fact that each element of prune parameters  $\alpha$  when set to zero reduces different number of BERT parameters, we multiply the  $\ell_1$  loss terms with the cost terms  $\beta$ ’s. For example,  $\beta_{a_i}$  is proportional to the number of model parameters that will be removed when an element of the prune parameter  $\alpha_{a_i}$  is set to zero. It is
critical to incorporate  $\beta$ 's. Their values are significantly different from each other. The  $\beta$  values are 1.0, 0.73, 0.093, 0.093, 0.0078 for  $a, h, k, v$  and  $f$  respectively.
After training the prunable BERT model for a fixed number of steps, we truncate the smallest prune parameters to zero, and remove the corresponding rows/columns from the BERT parameter matrices  $\{K_i,Q_i,V_i,P_i,D_i,G_i\}_{i\in [\ell ]}$ . Then we fine-tune the so obtained smaller schuBERT model.
Algorithm 1 summarizes our approach. If we want to reduce the number of parameters by a fraction  $\eta$ , we do so in  $T$  steps. In each step, we prune  $\eta / T$  fraction of parameters, and at the end of the step we fine-tune the network and repeat these steps  $T$  times. Though we have explained the algorithm in terms of  $\ell_1$  penalty on the prunable parameters, in our experiments we tried alternative sparsity inducing penalties as well-  $\ell_0$  regularization, and proximal gradient descent on prunable parameters.
$$
\begin{array}{l} \arg \min  _ {\left\{E, \left\{K _ {i}, Q _ {i}, V _ {i}, P _ {i}, D _ {i}, G _ {i} \right\} _ {i \in [ \ell ]} \right\}} \\ \mathcal {L} _ {\mathrm {M L M} + \mathrm {N S P}} (E, \{K _ {i}, Q _ {i}, V _ {i}, P _ {i}, D _ {i}, G _ {i} \} _ {i \in [ \ell ]}). \tag {1} \\ \end{array}
$$
$$
\arg \min _ {\{\alpha_ {h}, \{\alpha_ {a _ {i}}, \alpha_ {v _ {i}}, \alpha_ {k _ {i}}, \alpha_ {f _ {i}} \} _ {i \in [ \ell ]} \}}
$$
$$
\mathcal {L} _ {\mathrm {M L M + N S P}} (E, \{\widetilde {K _ {i}}, \widetilde {Q _ {i}}, \widetilde {V _ {i}}, \widetilde {P _ {i}}, \widetilde {D _ {i}}, \widetilde {G _ {i}} \} _ {i \in [ \ell ]})
$$
$$
\begin{array}{l} + \gamma \left\{\beta_ {h} \| \alpha_ {h} \| \right\} + \gamma \sum_ {i = 1} ^ {\ell} \left\{\beta_ {a _ {i}} \| \alpha_ {a _ {i}} \| + \beta_ {v _ {i}} \| \alpha_ {v _ {i}} \right\| \\ + \beta_ {k _ {i}} \| \alpha_ {k _ {i}} \| + \beta_ {f _ {i}} \| \alpha_ {f _ {i}} \| \}. \tag {2} \\ \end{array}
$$