Optimizing Transformers with Approximate Computing for Faster, Smaller and more Accurate NLP Models

1 INTRODUCTION . Transformer networks with hundreds of billions of parameters , such as T5 ( Raffel et al . ( 2019 ) ) , Megatron ( Shoeybi et al . ( 2019 ) ) , BERT ( Devlin et al . ( 2019 ) ) , GPT-2 ( Radford et al . ( 2019 ) ) and GPT-3 ( Brown et al . ( 2020 ) ) , have achieved state-of-the-art performance in several Natural Language Processing tasks . Model sizes are expected to grow further in the future as increasing the number of parameters has been shown to improve performance . For instance , increasing the number of parameters from 1.5B to 175B enabled a reduction in perplexity for Language Modelling ( Penn Treebank ) from 35.8 in GPT-2 to 20.5 in GPT-3 . This makes it computationally challenging to train Transformers as well as perform inference using them . The challenges associated with training these models are alleviated through the ( re- ) use of pre-trained models that are subsequently fine-tuned for different tasks . Consequently , these models incur a major one-time cost in computational resources , time and energy during the pre-training process , but the repeated fine-tuning for individual downstream tasks is performed at a considerably lower cost . However , performing inference using fine-tuned Transformer models continues to remain a challenge because of the large amount of storage and compute operations required . Prior research efforts have explored different techniques for improving the efficiency of Transformer inference . However , several of the proposed approaches either require training the network completely from scratch ( which is extremely compute and memory-intensive ) , or cause significant degradation in accuracy on the downstream task . In this work , we overcome these limitations by exploiting the transfer learning step in Transformers to produce individually optimized models for the different downstream tasks , using techniques that do not require training from scratch and maintain or improve accuracy levels . From the runtime and memory breakdown of Transformers ( Fig . 1 ) , we observe that the most timeconsuming and memory-intensive operations in a Transformer are the self-attention ( ATTN ) blocks , which are used to identify and form relationships between the different tokens in text , and the feedforward neural network blocks ( FFN blocks ) in the Transformer layers . These blocks together account for more than 85 % of the inference time ( and more than 75 % of the model ’ s parameters ) . We accordingly optimize the execution of these two components in our approach . The self-attention component dominates the execution time and memory size for longer context lengths as its operation scales quadratically in time and memory with sequence length . Some previous works ( Kitaev et al . ( 2020 ) , Ye et al . ( 2019 ) ) have addressed this issue , accelerating training and inference of Transformers when large context lengths are used . However , they suffer from significant overheads and slowdowns in applications with smaller context lengths , such as question answering , where questions and answers are usually short , in the order of a few hundred tokens . Our approach , on the other hand , performs well across context lengths , size of hidden layers , number of layers and other network characteristics . The pre-training of Transformer models with some initial objective ( most commonly predicting masked words in a large text corpus ) and the subsequent fine-tuning on a downstream task leads to highly over-parameterized models for many downstream tasks ( Michel et al . ( 2019 ) ) , providing ample opportunities for approximations . As these models grow larger , such opportunities are expected to increase even further . We observe that for a given downstream task , some parts of the pre-trained Transformer are more significant to obtain good accuracy , while other parts are less important or unimportant . In order to exploit this observation in a principled manner , we introduce a framework to introduce approximations while fine-tuning a pre-trained Transformer network , optimizing for either size , latency , or accuracy of the final network . We perform and apply significance analysis in a hierarchical manner , first pruning entire blocks , followed by attention heads , and finally pruning weight groups . We achieve further gains by also allowing elements that can not be pruned to be approximated by other techniques . We specifically apply two forms of approximations , depending on the element type . For weights , we utilize quantization . For the self-attention operation , we replace the scaled dot product attention mechanism with a novel sign matching-based attention mechanism . We summarize our main contributions as follows : • We introduce a framework for creating fine-tuned models from pre-trained Transformer models that are optimized for various metrics ( size , latency , accuracy ) . • We incorporate multiple heuristics in the framework , such as hierarchical processing , model-driven insights , and run-time based ordering of elements . • We propose a significance analysis technique to identify the importance of each element of the pre-trained Transformer for a given downstream task . We use this technique to prune entire blocks , attention heads , and weight groups and to guide the quantization of low-importance weights . • We propose a low-complexity attention mechanism , sign matching , in order to approximate dot product attention in the less significant attention layers . • Across a suite of different Transformer networks , including previously proposed optimized networks , we demonstrate that our techniques produce models that are up to 4× faster and up to 14× smaller ( with less than 0.5 % relative accuracy degradation ) , or up to 5.5 % more accurate with simultaneous size and latency improvements . 2 RELATED WORK . Given the effectiveness and popularity of Transformer models , several techniques have been proposed to overcome their computational and memory challenges , and to accelerate inference using these models . Most of these works directly pre-train efficient models from scratch . For example , DistilBERT ( Sanh et al . ( 2019 ) ) , MobileBERT ( Sun et al . ( 2020 ) ) and TinyBERT ( Jiao et al . ( 2019 ) ) use knowledge distillation to train smaller and faster networks using the original network as a teacher . LayerDrop ( Fan et al . ( 2020 ) ) randomly drops layers during pre-training , thereby enabling their dropping during inference . SchuBERT ( Khetan & Karnin ( 2020 ) ) learns the optimal sizes of the BERT elements during pre-training . Lite Transformer ( Wu et al . ( 2020 ) ) uses Long-Short Range Attention to speed up the self-attention operation , with different attention heads attending to local and global context . Depth-adaptive Transformer ( Elbayad et al . ( 2020 ) ) and DeeBERT ( Xin et al . ( 2020 ) ) modulate Transformer depth depending on the complexity of each input sample using gating functions that are trained along with the model . AlBERT ( Lan et al . ( 2020 ) ) uses factorized embeddings and cross-layer parameter sharing . These works are orthogonal to ours , as the models that they produce are still subsequently fine-tuned for downstream tasks . We demonstrate using DistilBERT , AlBERT and LayerDrop as examples that these optimized networks still have significant opportunities that our techniques can take advantage of . Other works ( including ours ) aim to improve the inference efficiency of Transformers using techniques that do not require training new models from scratch . Among these , PoWER-BERT ( Goyal et al . ( 2020 ) ) , which eliminates redundant word vectors from the model without removing any parameters , and Q8BERT ( Zafrir et al . ( 2019 ) ) , which quantizes all weights and activations in the model to 8-bit integers through the use of Quantization-Aware Training at fine-tuning time , are orthogonal and complementary to our work . Poor Man ’ s BERT ( Sajjad et al . ( 2020 ) ) evaluates several layer-dropping techniques that do not require re-training . Compared to layer-dropping techniques that do not require re-training , our techniques produce models that are up to 20 % more accurate at comparable inference speed , and this is especially true when working with highly optimized baselines such as Q8BERT . Our framework can also be adapted to satisfy a wide range of user constraints . 3 PRELIMINARIES . A Transformer ( Fig . 1 ) consists of an embedding layer , followed by multiple transformer layers stacked together , and a task-specific final layer . A transformer layer consists of the multi-headed self-attention operation ( ATTN block ) , followed by a feed-forward neural network ( FFN block ) with layer norm operations at the input and output of the layer . In this work , we define the elements of a Transformer to include different levels of granularity , i.e. , ATTN blocks , FFN blocks , Attention Heads and Weight Groups . We define Weight Groups only along dimensions that do not impact the shape of the output of the block when these groups are removed . The self-attention operation takes as input a sequence n of vectors X , and computes three matrices , Query = X × Wq , Key = X × Wk and Value = X × Wv . Then , the output of the self-attention operation is computed as Y = softmax ( ( Query × KeyT ) + attention mask ) × V alue . For auto-regressive models , tokens are not allowed to attend to future tokens . Hence , an attention mask is applied before the softmax operation , setting attention scores with future tokens to a very large negative number , which becomes zero after the softmax operation . This operation has multiple “ attention heads ” working in parallel on the input sequence , where each head has its own set of parameters to compute the query , key and value matrices . The independent attention outputs are concatenated and transformed into the expected output dimensions . The self-attention operation scales quadratically in time and memory with sequence length n since Query × KeyT has n2 entries . 4 DESIGN METHODOLOGY . We propose a framework for producing fine-tuned Transformer models that are optimized for a specific metric ( speed , model size , or accuracy ) . Fig . 2 presents an overview of the proposed framework . As shown in the figure , the inputs to the framework are a pre-trained Transformer model , the fine-tuning dataset , the goal of optimization ( speed , size or accuracy ) and acceptable accuracy loss ( when optimizing for speed or size ) . The framework has three major components : ( i ) a set of heuristics used to build an ordered queue of elements ( TransElements ) to be considered , ( ii ) a significance analysis method to identify insignificant elements in a pre-trained Transformer and ( iii ) a set of techniques to prune or approximate the insignificant elements . The framework proceeds in an iterative manner . That is , we first start with the original Transformer . We then remove an element from the TransElements queue , analyze its significance , and apply pruning/approximation techniques to the element . This results in new Transformer , where the element is replaced by the pruned or approximated version . This modified Transformer is then used as the baseline for the next iteration . After processing all of the identified elements , we fine-tune on the downstream task for the same number of epochs as the baseline model to obtain the final , optimized model . A detailed description of our methodology for approximating Transformers is presented in Fig . 2 and in Algorithm 4 . In the following subsections , we further describe our techniques for generating the ordered queue TransElements , followed by the significance analysis method , and finally the pruning and approximation techniques for different Transformer elements . TransElement Ordered Queue . In order to optimize a given model , we would ideally want to characterize the significance of each and every parameter in the model , rank them in order of importance , and finally prune/approximate only the least significant parameters , as in Molchanov et al . ( 2017 ) . However , Transformers have billions of parameters , making this process computationally infeasible . In addition , previously proposed techniques that can efficiently estimate the importance of each parameter , such as using Taylor expansion , are not applicable . This is because the { approximate , fine-tune , approximate } cycle does not work for Transformers during fine-tuning , since they very quickly overfit the training data for the downstream task ( usually within 5 epochs ) . We take advantage of hierarchical structure of Transformers and consider them in a hierarchical manner , ordered by increasing granularity . Specifically , we place entire FFN and ATTN blocks earlier in the queue , followed by heads , and finally weight groups . Through this ordering , we are able to quickly eliminate large numbers of parameters from further consideration , speeding up future iterations of the framework . For example , eliminating a single FFN block in the BERT-Base model removes 5.6 % of all parameters under consideration . To further reduce the number of elements under consideration , we also dynamically remove elements from the queue if they are encompassed by a high-importance block . For example , if a given ATTN block is determined to be of high importance , we remove all heads and weight groups within that block from the TransElement queue . Since the framework iterates through the entries of the TransElement queue sequentially , its efficacy is dependent on the ordering of the elements at each level of granularity . In order to minimize the run-time of the framework , we provide two additional heuristics to guide the ordering of elements . First , we use the unique linguistic properties captured by the different Transformer layers ( Jawahar et al . ( 2019 ) ) . These properties depend on both the Transformer and the downstream task under consideration , since different tasks require different types of linguistic knowledge . For example , top layers usually have low significance for Language Understanding tasks , since long-range dependency information is not required for most tasks ( for example , sentiment analysis requires only local context ) . Hence , we place the final layer at the front of the queue , and work our way backwards towards the first layer , since blocks in the final layers are more likely to be removed , thereby speeding up future iterations . Second , we use a run-time ( or parameter-count ) aware ordering of the TransElements , such the most time consuming blocks ( or blocks with the most parameters ) are likely to be removed earlier in the algorithm . For example , at large context lengths , we start with the ATTN blocks in all layers before moving on to the FFN blocks , and vice-versa at small context lengths . This has the dual benefit of producing highly optimized models for inference , as well as speeding up the significance analysis process by eliminating time-consuming blocks early and making further iterations faster . Algorithm 1 and Fig . 2 describe the process of creating the TransElement Queue . The utility of this framework and the heuristics used are discussed in Appendix C. Significance Analysis . To determine the significance of each Transformer element , we first finetune the original Transformer model for the given downstream task to obtain the baseline loss . We then use this baseline loss , along with the provided acceptable accuracy degradation , to generate a set of loss thresholds that determine whether a given element is of low importance and therefore can be pruned/approximated . This is a one-time step and performed globally for all elements in the TransElements queue . Then , for the element under consideration in each iteration of the framework , we compute the loss of the current Transformer model with the element removed . We then compare this loss to the thresholds determined above . The exact thresholds used are dependent on the optimization metric : speed , size , or accuracy . If we are optimizing the network for speed or size , we prune the element under consideration if the training/validation loss upon removing it from the Transformer is less than the pruning threshold . If we are optimizing for accuracy , we prune the element only if the training/validation loss when it is removed is less than the minimum loss seen thus far during the optimization process , since the goal is to find a model with minimum loss . Similarly , we apply approximations if the loss with the element removed from the Transformer is greater than the pruning threshold but lower than the approximation threshold . Algorithm 2 describes Significance Analysis . Pruning and Approximating . As evident from Section 3 , the structure and functionality of ATTN blocks differ significantly from that of FFN blocks in a Transformer . We accordingly adopt different strategies for approximating them , as described below . But pruning an entire ATTN or FFN block is effectively the same as it simply involves using the skip connection to bypass that block . The pruning strategies for the FFN and ATTN blocks are illustrated in Fig . 4 and Fig . 5 . Pruning Weight Groups within approximable FFN Blocks . Consider an approximable FFN block that performs the transformation Rn×d × Rd×y → Rn×y , with weight groups defined along the d dimension ( ( d/W ) weight groups of ( W ) weights each , where W is a hyperparameter that defines the granularity of approximations ) . When optimizing models for accuracy , we remove weight groups only if doing so results in a reduction in the model loss . When optimizing for size , we remove weight groups that maintain loss within the pruning threshold when removed . When optimizing for speed , however , removing weight groups with low significance from arbitrary locations does not help , since it introduces unstructured sparsity in the weight matrix that can be difficult to exploit to achieve speedups . Instead , we impose structure on our pruning . Specifically , we use a “ greedy shrinking ” algorithm that finds the largest number of weight groups that can be removed while maintaining loss below the threshold , such that the weight groups that remain in the model form a contiguous block . We first start from the bottom ( weight group 0 ) , work our way up and remove as many weight groups as possible while staying within the loss threshold . We then start from the top ( weight group d/W ) , work our way down and remove as many weight groups as possible while staying within the loss threshold . When this process is completed , the weight groups that remain form a contiguous dense block , enabling speedups on all hardware platforms . Since weight groups are removed along the “ hidden ” dimension d , our methods do not change the shape of the output of this block ; instead , we are simply “ shrinking ” the effective hidden dimension size through structured pruning . Quantizing Weight Groups within approximable FFN and ATTN Blocks . When optimizing the Transformer for size , we also quantize weight values within weight groups for which the loss lies between the pruning and approximation thresholds . We use uniform quantization with QuantizationAware Training proposed in Q8BERT ( Zafrir et al . ( 2019 ) ) within our hierarchical framework to quantize insignificant weight groups to lower precisions . This reduces the memory requirements of those weight groups but does not improve the execution time as the computations are still performed at the baseline precision . Pruning ATTN heads and Weight Groups within approximable ATTN Blocks . We divide the multi-headed self-attention operation into two main steps . In the first step , we compute the Query , Key and Value matrices by multiplying the input to this layer with the corresponding weight matrices ( Rn×d × Rd×y → Rn×y ) , and then reshape them into multiple attention heads ( Rn×y → Rn×h× ( y/h ) ) . Our approach to pruning this step is exactly the same as for the FFN blocks , where we iteratively prune weight groups along the d dimension using our shrinking algorithm . In the second step , we compute the “ attention output “ as Y = softmax ( ( Query×KeyT ) + attention mask ) × V alue . To optimize this step , we apply two techniques . Firstly , we identify insignificant attention heads , and prune them from the model . However , removing attention heads changes the shape of the output of this layer . We overcome this by keeping track of the pruned heads , and padding the output with zeros in the corresponding locations . In spite of this overhead , we still manage to achieve significant speedup from this approximation technique since pruning heads makes multiple downstream operations ( computing the attention scores , applying softmax to the attention scores , and computing the final score ) considerably faster . Therefore , we do not use our greedy shrinking method , but rather rely on unstructured pruning as it allows for greater pruning which further benefits the downstream operations . Secondly , we dynamically reduce the size of the key and value matrices by pruning weight groups from the same location along the n dimension in both matrices , based on sign matches with the query vectors . This again makes multiple downstream operations considerably faster and does not change the shape of the output of the pruned block . Approximating self-attention within approximable ATTN Blocks . We observe that the “ attention scores ” matrix is highly sparse , especially after the softmax operation . This sparsity implies that most of the dot products between the query and the key are unnecessary . Thus , we would ideally like to perform the attention operations for the query vectors that give highest dot-product with each key vector efficiently without explicitly performing all of the dot products . To this end , we propose replacing the O ( n2 ) dot product-based attention mechanism with a linear-time sign-matching-based mechanism in approximable ATTN blocks . Sign-matching attention ( SM ) is based on the idea that key vectors whose signs match with the largest number of query vectors will have high dot-products with maximum number of query vectors . However , it is expensive to compute a sign match for all pairs of query-key vectors , as this will grow quadratically . Instead , we employ a low-cost approximation . For each column of the query matrix , we identify if more number of vectors will have a positive or negative number in that column . This becomes the representative sign in that column for all the query vectors . Each key vector is then scored by how well the sign of each of its elements matches with the sign of the representative query vector by computing the Hamming distance between the two sign vectors . This score is used to select the top K key vectors . As a result , we reduce the number of computations required to score each key vector ( and the overall complexity ) from O ( n2 ) to O ( n ) . Sign matching is illustrated in Fig . 2 , and explained in detail in Appendix B . As this approximation does not increase the accuracy of the models nor decrease the number of parameters , they are only applied when optimizing the fine-tuned models for speed .

Summary: This paper presents a method for improving a fine-turned Transformer in terms of a specific metric such as size, speed, or accuracy. The candidates of removed elements are considered hierarchically with some heuristics and are evaluated in terms of training and validation loss to determine whether they should actually be removed from the model. The authors apply their method to several state-of-the-art Transformer models and show that they can produce fast and compact models without losing much accuracy.