Compressing Large-Scale Transformer-Based Models: A Case Study on BERT (2021)

Abstract: Pre-trained Transformer-based models have achieved state-of-the-art performance for various Natural Language Processing (NLP) tasks. However, these models often have billions of parameters, and thus are too resource- hungry and computation-intensive to suit low- capability devices or applications with strict latency requirements. One potential remedy for this is model compression, which has attracted considerable research attention. Here, we summarize the research in compressing Transformers, focusing on the especially popular BERT model. In particular, we survey the state of the art in compression for BERT, we clarify the current best practices for compressing large-scale Transformer models, and we provide insights into the workings of various methods. Our categorization and analysis also shed light on promising future research directions for achieving lightweight, accurate, and generic NLP models.

## introduction sentiment analysis, paraphrase detection, machine reading comprehension, question answering, text summarization-all these natural language processing (nlp) tasks benefit from pre-training a large-scale generic model on an enormous corpus such as a wikipedia dump and/or a book collection, and then fine-tuning for specific downstream tasks, as shown in figure 1 . earlier solutions following this methodology used recurrent neural networks (rnns) as the base model, for example, ulmfit @xcite and elmo @xcite , but more recent methods use the transformer architecture @xcite , which relies heavily on the attention mechanism. popular pre-trained transformers include bert @xcite , gpt-2 @xcite , xlnet @xcite , megatronlm @xcite , turing-nlg @xcite , t5 @xcite , and gpt-3 @xcite . these transformers are-for example, bert, when first released, improved the state of the art for eleven nlp tasks by sizable margins @xcite . however, transformers are also bulky and resource-hungry: for instance, gpt-3 @xcite , a recent large-scale transformer, has over 175 billion parameters. models of this size incur high memory consumption, computational overhead, and energy. the problem is exacerbated when we consider devices with lower capacity (e.g., smartphones), and applications with strict latency constraints, (e.g., interactive chatbots). to put things in perspective, a single training run for gpt-3 @xcite , one of the most powerful and heaviest transformer-based models, trained on a total of 300 billion tokens, costs well above 12 million usd @xcite . moreover, fine-tuning or even inference with such a model on a downstream task cannot be done on a gpu with 32gb memory, which is the capacity of tesla v100, one of the most advanced data center gpus. instead it requires access to high-performance gpu or multi-core cpu clusters, which often means a need to access cloud computing with high computation density, such as the google cloud platform (gcp), microsoft azure, amazon web services (aws), and so forth, and results in a high monetary cost @xcite . one way to address this problem is through model compression, an intricate part of deep learning that has attracted attention from both researchers and practitioners. a recent study by @xcite highlights the importance of first training over-parameterized models and then compressing them, instead of directly training smaller models, to reduce the performance errors. although most methods in model compression were originally proposed for convolutional neural networks (cnns) (pruning, quantization, knowledge distillation, etc.) @xcite , many ideas are directly applicable to transformers. there are also methods designed specifically for transformers (e.g., attention head pruning, attention decomposition, replacing transformer blocks with an rnn or a cnn), which we will discuss in section 3. unlike cnns, a transformer model has a relatively complex architecture consisting of multiple parts such as embedding layers, self-attention, and feed-forward layers (details introduced in section 2). thus, the effectiveness of different compression methods can vary when applied to different parts of a transformer model. several recent surveys have focused on pretrained representations and large-scale transformerbased models @xcite @xcite . however, to the best of our knowledge, no comprehensive, systematic study has compared the effectiveness of different model compression techniques on transformerbased large-scale nlp models, even though a variety of approaches for compressing such models have been proposed. motivated by this, here we offer a thorough and in-depth comparative study on compressing transformer-based nlp models, with a special focus on the widely used bert @xcite . although the compression methods discussed here can be extended to transformer-based decoders and multilingual transformer models, we restrict our discussion to bert in order to be able to provide more detailed insights into the various methods that we compare. our study is timely, since (i) the use of transformer-based bert-like models has grown dramatically, as demonstrated by current leaders of various nlp tasks such as language understanding @xcite , machine reading comprehension @xcite , machine translation @xcite , summarization @xcite , and so on; (ii) many researchers are left behind as they do not have expensive gpus (or a multi-gpu setup) with a large amount of gpu memory, and thus cannot fine-tune and use the large bert model for relevant downstream tasks; and (iii) aipowered devices such as smartphones would benefit tremendously from an on-board bert-like model, but do not have the capability to run it. in addition to summarizing existing techniques and best practices for bert compression, we point out several promising future directions of research for compressing large-scale transformer-based models. bert decomposes the input sentence(s) into wordpiece tokens @xcite . specifically, wordpiece tokenization helps improve the representation of the input vocabulary and reduce its size, by segmenting complex words into subwords. these subwords can even form new words not seen in the training samples, thus making the model more robust to out-of-vocabulary (oov) words. bert further inserts a classification token ([cls]) before the input tokens, and the output corresponding to this token is used for classification tasks that target the entire input. for sentence pair tasks, the two sentences are packed together by inserting a further separator token ([sep]) between them. bert represents each wordpiece token with three vectors, namely, its token, segment, and position embeddings. these embeddings are summed together and then passed through the main body of the model (i.e., the transformer backbone), which produces the output representations that are fed into the final, application-dependent layer (e.g., a classifier for sentiment analysis). as shown in figure 2 , the transformer backbone consists of multiple stacked encoder units, each with two major sub-units: a self-attention sub-unit and a feed forward network (ffn) sub-unit, both with residual connections. each self-attention sub-unit consists of a multi-head self-attention layer, and fully connected layers before and after it. an ffn sub-unit exclusively contains fully connected layers. the architecture of bert can be specified using the following three hyper-parameters: number of encoder units (l), size of the embedding vector (h), and number of attention heads in each self-attention layer (a). l and h determine the depth and the width of the model, whereas a is an internal hyper-parameter that affects the number of contextual relations that each encoder can focus on. the authors of bert provided two pre-trained models: we conducted various experiments with the bert base model by running inference on a sen-tence of length 256, and then we collected the results in figure 3 . the top graph in the figure compares the model size as well as the theoretical computational requirements (measured in millions of flops) of different parts of the model. the bottom two graphs track the model's run-time memory consumption as well as the inference latency on two representative hardware setups. we conducted our experiments using nvidia titan x gpu with 12gb of video ram and intel xeon e5-1620 cpu with 32 gb of system memory, which is a commonly used server or workstation configuration. all data was collected using the pytorch profiling tool. clearly, the parts consuming the most memory in terms of model size and executing the highest number of flops are the ffn sub-units. the embedding layer is also a substantial part of the model size, due to the large vector size (h) used to represent each embedding vector. note that it has zero flops, since it is a lookup table that involves no arithmetic computations at inference time. for the self-attention sub-units, we further break down the costs into multi-head self-attention layers and the linear (i.e., fully connected) layers before and after them. the multi-head selfattention does not have any learnable parameters; however, its computational cost is non-zero due to the dot products and the softmax operations. the linear layers surrounding each attention layer incur additional memory and computational overhead, though it is relatively small compared to the ffn sub-units. note that the input to the attention layer is divided among various heads, and thus each head operates in a lower-dimensional space (h/a). the linear layer before attention is roughly three times the size of that after it, since each attention has three inputs (key, value, and query) and only one output. the theoretical computational overhead may differ from the actual inference cost at run-time, which depends on the hardware that the model runs on. as expected, when running the model on a gpu, the total run-time memory includes memory both on the gpu side and on the cpu side, and it is greater than for a model running solely on a cpu due to duplicate tensors present on both devices for faster processing on a gpu. the most notable difference between the theoretical analysis and the run-time measurements on a gpu is that the multi-head self-attention layers are significantly more costly in practice than in theory. this is because the operations in these layers are complex, and are implemented as several matrix transformations followed by a matrix multiplication and a softmax. moreover, gpus are designed to accelerate certain operations, and thus can implement linear layers faster and more efficiently than the more complex attention layers. when we compare the run-time performance on a cpu, where the hardware is not specialized for linear layer operations, the inference time as well as the memory consumption of all the linear layers shoots up more compared to the multi-head self-attention. thus, on a cpu, the behavior of runtime performance is similar to that of theoretical computations. the total execution time for a single example on a gpu (57.1 ms) is far superior as compared to a cpu (750.9 ms), as expected. the execution time of the embedding layer is largely independent of the hardware on which the model is executed (since it is just a table lookup) and it is relatively small compared to the other layers. the ffn sub-units are the bottleneck of the whole model, which is consistent with the results from the theoretical analysis. ## compression methods because of bert's complex architecture, no existing compression method has focused on every single aspect of the model: self-attention, linear layers, embedding size, model depth, and so forth. instead, each compression method applies to certain components of bert. below, we consider the compression methods that offer model size reduction and speedup at inference time, rather than during the training procedure. ## effectiveness of compression methods in this section, we compare the performance of several bert compression techniques based on their model size and speedup, as well as their accuracy or f1 score on various nlp tasks. we chose work whose results are either on the pareto frontier @xcite or representative for each compression technique mentioned in the previous section. ## open issues and research directions from our analysis and comparison, we conclude that traditional model compression methods such as quantization and pruning are beneficial for bert. techniques specific to bert also yield competitive results, for example, variants of knowledge distillation and methods that reduce the number of architectural hyper-parameters. such methods also offer insights into bert's workings and the importance of various layers in its architecture. we see multiple avenues for future research: