Specialized Transformers: Faster, Smaller and more Accurate NLP Models

1 INTRODUCTION . Transformers models such as BERT ( Devlin et al. , 2019 ) , GPT-2 ( Radford et al. , 2019 ) and GPT3 ( Brown et al. , 2020 ) have revolutionized the field of Natural Language Processing ( NLP ) , greatly advancing the state-of-the-art in many NLP tasks . Models that achieve good performance on these tasks are of high practical significance , finding their place in commercial applications such as social media monitoring ( sentiment analysis ) , AI chat assistants ( question answering ) , automated summarization tools ( analyzing sentence similarity ) , etc . Therefore , there is a strong interest in creating more accurate and efficient Transformer models for these tasks . Transformers are first pre-trained on very large datasets , and subsequently fine-tuned for different downstream tasks . However , the current method of pre-training and fine-tuning Transformers has two major drawbacks . First , fine-tuning the pre-trained Transformers leads to models that are highly over-parameterized models for the downstream tasks , especially since many of these tasks have very limited training data . This could lead to unstable models ( Dodge et al. , 2020 ) with sub-optimal generalization ability ( Michel et al. , 2019 ) . Second , these large fine-tuned models present high computation and storage requirements for inference . This problem is further exacerbated by the trend towards larger and more accurate models over time . For instance , increasing the number of parameters from 1.5B to 175B enabled a reduction in perplexity for Language Modelling ( on the Penn Treebank dataset ) from 35.8 in GPT-2 ( 2019 ) to 20.5 in GPT-3 ( 2020 ) . In this work , we address both these challenges , taking advantage of the over-parameterized nature of pre-trained models to create individually Specialized models for the different downstream tasks that are smaller , faster and more accurate than the conventional fine-tuned models . Prior research efforts have explored approximation techniques , such as quantization , pruning and distillation , for improving the inference efficiency of various classes of neural networks . However , these techniques invariably involve a trade-off between accuracy and efficiency . In addition , the vast majority of these techniques require re-training or additional fine-tuning . This becomes especially problematic for Transformers , since these large models require significant time and energy to train and fine-tune . Fine-tuning is usually performed on a very small data set , and as a result , conventional iterative prune-and-retrain methods quickly end up overfitting when applied to Transformer models . In contrast , our Specialization framework utilizes the unique characteristics of Transformer training and deployment to enable substantial gain in both accuracy as well as efficiency . Our framework also does not require any additional re-training or fine-tuning , and can be applied in a plug-and-play manner to any Transformer model that is fine-tuned for any downstream task . During pre-training , Transformers capture rich linguistic knowledge and gain a deep understanding of the structure of the target language . Fine-tuning refines this knowledge for a specific downstream task by training a task-specific final layer . We observe that , due to the nature of their design process , Transformer models are not only over-parameterized but also contain parts that are , in fact , harmful to performance of downstream tasks . In order to exploit this observation in a principled manner , we introduce a framework to identify and prune the harmful elements of a Transformer ( parameters , grouped at different levels of granularity i.e. , self-attention blocks , feed-forward neural network blocks , attention heads and neurons ) , with the goal of maximizing accuracy on the downstream task . In contrast with prior pruning methods that prune elements with little-or-no impact on the network output , the proposed method prunes elements that have a considerable impact on the output , leading to the highest positive impact on accuracy . In order to reduce the large pruning space , we analyze the different elements of the fine-tuned Transformer in a hierarchical manner , starting with entire self-attention or feed-forward neural network blocks , followed by attention heads , and neurons , and prune the harmful elements . The core of the Transformer is self-attention , where each token in the input builds its representation based on the extent of attention it places on all the other tokens . However , we observe that in some cases , restricting the attention span of each token to only focus on the relevant tokens ( in certain layers ) leads to better information flow inside the model . Hence , our framework is also equipped with the ability to identify the appropriate layers and replace the ” soft ” self-attention with ” hard ” self-attention in these layers . We also introduce Transformer-specific heuristics to minimize the run-time of our framework , thereby enabling it to scale to large Transformer models . We summarize our main contributions as follows : • We introduce a Specialization framework that optimizes Transformer models for specific downstream tasks through the use of accuracy-driven pruning and selective hard attention . • We incorporate multiple heuristics in the framework , such as hierarchical processing , model-driven insights , and run-time based ordering of elements , in order to minimize the overheads . • We propose a significance analysis technique to identify the importance of each element of the fine-tuned Transformer for a given downstream task . We use this technique to prune elements that are harmful to performance on the downstream task . • We propose the selective replacement of the ” soft ” self-attention with hard attention in the appropriate layers , helping the model focus only on the relevant parts on the input to build better representations . • Across a suite of different Transformer networks , we demonstrate that Specialized models are consistently more accurate and stable , while also being significantly faster and smaller than their conventional fine-tuned non-Specialized counterparts . 2 RELATED WORK . Task-agnostic optimizations . 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 . A vast majority of these works introduce task-agnostic optimizations , using popular approximation techniques such as knowledge distillation ( Sanh et al. , 2019 ; Sun et al. , 2020 ; Wang et al. , 2020 ) , early exit/ depth modulation ( Elbayad et al. , 2020 ; Xin et al. , 2020 ; Zhou et al. , 2020 ) , quantization ( Zafrir et al. , 2019 ) , attention head pruning ( Zhang et al. , 2020 ) and parameter sharing ( Lan et al. , 2020 ) . In addition , Fan et al . ( 2020 ) randomly drop layers during pre-training , thereby enabling their dropping during inference ; Khetan & Karnin ( 2020 ) learn the optimal sizes of the BERT elements during pre-training , and Wu et al . ( 2020 ) use Long-Short Range Attention to speed up the self-attention operation . Using DistilBERT and Q8BERT as examples , we demonstrate that our techniques are complementary to these works , and can be applied while fine-tuning for a specific task to create accurate ultra-efficient models for specific downstream tasks . Task-specific optimizations . Hou et al . ( 2020 ) , Shen et al . ( 2020 ) , Jiao et al . ( 2019 ) , Wang et al . ( 2020 ) , Lagunas et al . ( 2021 ) and Sajjad et al . ( 2020 ) introduce task-specific optimizations , but the gain in efficiency comes at the cost of degradation in accuracy on the downstream task . In contrast , our framework improves both accuracy as well as efficiency , and hence appeals to a wider range of users . We also demonstrate that task-specific optimizations are more effective when applied to Specialized models compared to conventional fine-tuned models . Pruning techniques . Finally , Structured Pruning has been applied to various classes of neural networks ( Anwar et al. , 2017 ; Molchanov et al. , 2017 ) , and greedy pruning strategies have also been explored to identify weights and parameters that the output is least sensitive to ( Zhuang et al. , 2018 ; Ye et al. , 2020 ) . In contrast , our method is designed to identify and prune parameters that have most detrimental effect on the output . 3 METHODOLOGY TO SPECIALIZE TRANSFORMERS . We propose a framework for producing Specialized Transformer models that are optimized for a specific downstream task , illustrated in Algorithm 1 . Our framework performs two main optimizations : ( 1 ) It identifies and prunes elements that hinder performance on the downstream task at hand . ( 2 ) It selectively replaces soft self-attention with hard self-attention to help the model focus only on the relevant parts of the input . 3.1 ACCURACY-DRIVEN PRUNING . The problem of identifying an optimal set of elements to prune is challenging , and this is especially true for Transformers . 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 only the least significant parameters . 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 limited training data for the downstream tasks ( usually within 5 epochs ) . We address both these issues through the use of a hierarchical greedy algorithm that does not require any additional training or fine-tuning . To determine the significance of each Transformer element , we first fine-tune the original Transformer model for the given downstream task to obtain the baseline loss . 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 prune the element under consideration if the 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 . Also , in order to prevent overfitting to the validation set , we introduce a generalization constraint in addition to the aforementioned loss condition . This constraint ensures that an element is pruned only if it decreases ( or retains ) the loss of at least a certain number ( N ) of samples in the validation set over the current best solution ( computed by num samples helped function in Alg . 1 ) , where N is the number of elements in the validation dataset whose loss is less than the average loss of the misclassified samples ( we consider samples whose loss is greater than the average loss of the misclassified samples to be outliers ) . Therefore , elements are pruned only if a vast majority of the samples in the validation set benefit from their removal , resulting in improved generalization performance . If the loss with the element removed is greater than the minimum loss Algorithm 1 : Transformer Specialization Input : Fine-tuned ( for the given downstream task ) Transformer T , Validation set D Output : Specialized Transformer for the given downstream task T Function analyze element ( element E ) : Tpruned = T − E New Loss = Evaluate ( Tpruned , D ) if New Loss < Min Loss and num samples helped > N then Min Loss = New Loss T = Tpruned Function num samples helped : computes the number of samples in the validation set that benefit from the pruning of an element ; must be > N for an element to be pruned Baseline Loss = Evaluate ( T , D ) Min Loss = Baseline Loss Q = Order elements for inspection ( T , D ) for each layer L in T do Replace soft self-attention in L with hard self-attention New Loss = Evaluate ( T , D ) if New Loss < Min Loss and num samples helped > N then Min Loss = New Loss else Restore soft self-attention in L while Q is not empty do TrialElement = Q.pop ( ) analyze element ( TrialElement ) if TrialElement has not been pruned from T and New Loss < Baseline Loss then if TrialElement is an attention block then for each attention head h in TrialElement do analyze element ( h ) else if TrialElement is a feed-forward block then for each neuron w in TrialElement do analyze element ( w ) return T seen so far but less than the baseline loss , we inspect the element at a finer granularity , and prune only parts of the element that hinder performance ( rather than pruning the entire element ) . Hierarchical processing of elements . It is computationally prohibitive to analyze every single parameter in large Transformers using the method described in Alg 1 . Since the framework iterates through the entries of the queue sequentially , its efficacy is dependent on both the total number of elements under consideration , and the time required to analyze each element . We take advantage of the inherently hierarchical structure of Transformers and consider the elements in a hierarchical manner , ordered by increasing granularity . Specifically , we analyze entire feed-forward and selfattention blocks first , and inspect them at finer granularity ( attention heads and neurons ) only when required . Through this ordering , we are able to quickly eliminate large numbers of parameters from further consideration . In addition , due to the over-parameterized nature of Transformers , it is likely that time-consuming blocks are pruned from the Transformer earlier in the process , thereby speeding up future iterations of the framework . For example , eliminating a single feed-forward block in the BERT-Base model removes 5.6 % of all parameters under consideration , and speeds up future iterations by 1.15× . To further reduce the number of elements under consideration , we also dynamically remove elements if they are encompassed by a high-importance block . For example , if a given self-attention block is determined to be of high importance ( the validation loss with the block removed is greater than the baseline loss ) , we remove all heads within that block from further consideration . Creating an ordered queue of elements for inspection . Since our framework performs greedy pruning of highly over-parameterized models , it is essential to know where the harmful elements are likely to be . Some elements may appear to be harmful for the downstream task ( especially in early iterations of the framework ) , but this often ends up being an artifact of over-parameterization . As a result , when elements are not carefully ordered for inspection , our framework lands in local minima of the validation loss function , leading to inefficient models with sub-optimal generalization ability . Our solution to this problem utilizes the unique linguistic properties captured by the different Transformer layers ( Jawahar et al. , 2019 ) to guide the ordering of elements for inspection . For example , it was found that BERT captures phrase-level information in the lower layers , mapping related tokens together . The lower layers also capture surface features , while the middle layers capture syntactic features and higher layers capture semantic features . It was also observed that BERT requires deeper layers only when long-range dependency information is required . Different tasks require different types of linguistic knowledge . For example , sentiment analysis requires only local context , and long-range information often ends up confusing the model , since sentiments often change rapidly ; it is also unlikely that syntactic and semantic information are needed . 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 hinder performance on sentiment analysis . This ordering of elements ensures that elements that are pruned early in our framework ( when the model is most over-parameterized ) do not lead the system into bad local minima .

Summary: This paper proposes a new framework for pruning a pretrained Transformer on downstream tasks for better performance and efficiency. The resulted model is 2.5x faster and 3.2x smaller than a fine-tuned model. The authors test the effectiveness on GLUE and SQUAD with BERT-base, Q8BERT, DistilBERT and XLNet.