schu BERT : Optimizing Elements of BERT (2020)

Abstract: Transformers have gradually become a key component for many state-of-the-art natural language representation models. A recent Transformer based model- BERTachieved 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.

## introduction transformer @xcite based models have achieved state-of-the-art performance for many natural language processing tasks (dai and @xcite @xcite . these include machine translation @xcite , question-answering tasks @xcite , natural language inference @xcite and semantic role labeling @xcite . a recent transformer based model bert @xcite 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 @xcite . following the seminal results obtained by the bert model, several follow up studies explored methods for improving them further. xlnet @xcite adds autoregressive capabilities to bert, improving its quality, though at the cost of additional compute requirements. roberta @xcite 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. @xcite 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). al-bert @xcite 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 @xcite 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 both the pre-training loss and the number of model parameters. our experiments result in the following findings: 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 @xcite provides an ablation study, 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. ## related works there is a vast literature on pruning trained neural networks. starting with the classical works @xcite ; @xcite in the early 90's to the recent works @xcite , pruning deep neural networks has received a lot of attention. there have been two orthogonal ap-proaches in pruning networks: structured pruning @xcite and unstructured pruning @xcite . structured pruning gives smaller architecture whereas unstructured pruning gives sparse model parameters. in natural language processing, @xcite explored structured pruning in feed-forward language models. @xcite and @xcite 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 @xcite @xcite . quantization is another technique to reduce the model size. this is done by quantizing the model parameters to binary @xcite , ternary @xcite , or 4 or 8 bits per parameter @xcite . recently published distilbert @xcite 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 ## 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 ℓ. 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-sentenceprediction (nsp) loss. for further details of the loss bert-base number of encoder layers ℓ 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 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 nonidentical layers. this implies that a generalized 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 × a 1 × 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 × a 1 × 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 @xcite . 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 × 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 × h and g 1 consists of a parameter matrix of dimension h × 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. 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 × a 1 × h. the value vector can be of different dimension v 1 . therefore the value tensor v 1 should be of dimension v 1 × a 1 × 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 @xcite and roberta @xcite , we set the wordpiece embedding size e equal to the hidden layer size h, i.e. e ≡ h. however, factorization of the embedding matrix can be incorporated as demonstrated in albert @xcite . ## optimization method we optimize bert design dimensions listed in for each design dimension that we seek to optimize, we introduce a prune-parameter vector α 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 α f 1 ∈ 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 our approach is to optimally find which individual elements of prunable parameters {α h , {α a i , α v i , α k i , α f i } @xmath1 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 @xcite . 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 } @xmath2. next, we introduce the prunable parameters given in 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 } @xmath3. 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 η, we do so in t steps. in each step, we prune η/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 ℓ 1 penalty on the prunable parameters, in our experiments we tried alternative sparsity inducing penalties as well-ℓ 0 regularization, and proximal gradient descent on prunable parameters. ## experimental results in this section, we present our experimental results. we apply algorithm 1 on bert base. for pre-training bert base we use mxnet based gluon-nlp repository that uses the hyper-parameters suggested in the original bert paper. besides pretraining, our algorithm has three hyper-parameters: regularization coefficient γ, learning rate for prunable parameters, and the number of steps for regularizing prune parameters-equation (2). we run hyper-parameter optimization on these parameters to get the best results. for regularization loss (2), we use the same training data that we use for pre-training, bookscorpus (800m words) and english wikipedia (2,500m words). however, we run the regularization step for 1/1000th steps as used for pre-training. we finet-une the pruned bert algorithm 1 pruning transformers input: a transformer model, minimization objective (flops/params/latency), target fraction η, number of iterations t . output: a optimally pruned transformer model. pre-training: train the network using loss equation (1). repeat t times: we provide accuracy results for schubert on the following downstream tasks-question answering datasets-squad v1.1, squad v2.0; and glue datasets -mnli, mrpc, sst-2 and rte. for these downstream tasks, we use the fine-tuning hyper-parameters as suggested in the bert paper. we create six schuberts by pruning one or all of the design dimensions. accuracy of the downstream tasks on these schuberts are given in tables 6-13. the bert base has 108 million parameters. the schubert sizes 88, 66, 43 million are chosen to match the number of parameters in bert with ℓ ∈ {9, 6, 3} layers. we use schubert-x notation for @xmath4, f, a} to denote a schubert obtained by only pruning h-hidden size, f -filter size of feed-forward, anumber of attention heads respectively. we use schubert-all to denote the case when all the design dimensions-h, f, a, k, v, except ℓ are pruned. we compare our results with original bert base, and by varying its number of encoder layers ℓ ∈ {12, 9, 6, 3}. we denote these results by bert-ℓ. since albert reduces parameters by factorizing the embedding matrix, we denote its results by 89.9/82.9 89.9/82.9 89.9/82.9 80.1/77.8 80.1/77.8 80.1/77.8 82.9 -91.5 -- tall and narrow bert. when we aim to obtain a computationally lighter model, using a 'tall and narrow' architecture provides better performance than a 'wide and shallow' architecture. our results in tables 8 , 11 , 13 demonstrate that schu-bert with ℓ = 12 encoder layers significantly outperforms bert with ℓ ∈ {9, 6, 3} layers for the same number of parameters.