Token and Head Adaptive Transformers for Efficient Natural Language Processing (2022)

Abstract: While pre-trained language models like BERT have achieved impressive results on various natural language processing tasks, deploying them on resource-restricted devices is challenging due to their intensive computational cost and memory footprint. Previous approaches mainly focused on training smaller versions of a BERT model with competitive accuracy under limited computational resources. In this paper, we extend Length Adaptive Transformer and propose to design Token and Head Adaptive Transformer, which can compress and accelerate various BERT-based models via simple fine-tuning. We train a transformer with a progressive token and head pruning scheme, eliminating a large number of redundant tokens and attention heads in the later layers. Then, we conduct a multi-objective evolutionary search with the overall number of floating point operations (FLOPs) as its efficiency constraint to find joint token and head pruning strategies that maximize accuracy and efficiency under various computational budgets. Empirical studies show that a large portion of tokens and attention heads could be pruned while achieving superior performance compared to the baseline BERT-based models and Length Adaptive Transformers in various downstream NLP tasks. MobileBERT trained with our joint token and head pruning scheme achieves a GLUE score of 83.0, which is 1.4 higher than Length Adaptive Transformer and 2.9 higher than the original model.

## introduction the transformer @xcite has become the dominant neural architecture used in natural language processing. especially, pretrained models using the transformer architecture as their backbone, e.g., bert @xcite , roberta @xcite , and gpt-3 @xcite , have shown significant accuracy im-provement across various natural language processing tasks. however, the amount of computation required for these large neural networks is enormous, which leads to a substantial increase in energy consumption and carbon footprint. for example, an article points out that in order to train a gpt-3 networks using standard nvidia gpus, it would roughly take 190,000 kwh of energy which is equivalent to producing around 85,000 kg co 2 considering america's average carbon emission intensity @xcite . despite their remarkable performance, these power hungry models pose a great hindrance along the pathway to sustainable systems. hence, it is crucial to make the large-scale networks efficient and sustainable to be deployed on various hardware with different computational budgets and adopted for real-world applications. there have been efforts to improve the efficiency of the large language models (see section 6 for a discussion in detail.) one of the effective approaches has focused on removing redundant model parameters. recent works @xcite show that only a small subset of heads are important and mainly attributed to the final decision making and some heads could even be pruned at test time to improve efficiency. other works successfully improve the efficiency of bert-based models by pruning out word token vectors. length adaptive transformer @xcite based on power-bert @xcite demonstrates that training an adaptive transformer with progressive word token pruning improves accuracyefficiency trade-off that can satisfy any target computational budget. in this paper, we focus on developing a framework capable of training an adaptive transformer that jointly eliminates redundant attention heads and word tokens to maximize accuracy and minimize required computational resources. we perform a multi-objective evolutionary search to find a full pareto frontier of joint token and attention-head pruning schemes that provides optimal accuracy-efficiency trade-offs given any computational budget. we applied our framework to bert-based models with different sizes. our progressive head pruning scheme combined with word token pruning allows heavy head pruning in the last layers of bert, mobilebert, and distilbert @xcite without any accuracy loss. furthermore, the final models trained with our framework are robust to prune even more tokens along with attention heads. this results in superior accuracy and latency tradeoff than simply pruning word tokens. empirical evaluations show that mobilebert trained with our joint token and head pruning scheme achieves a glue score of 83.0, which is 1.4 higher than length adaptive transformer and 2.9 higher than the original model. our adaptive bert and dis-tilbert also achieve higher accuracy with less computational cost compared to length adaptive transformers. on squadv1.1 @xcite question answering task, token and head adaptive transformers (that) obtain higher dev f1 scores with less computational resources compared to length adaptive transformers and the original bert-based models. ## background in this section, we present the main building blocks of our framework. we review multi-head attention, the essential mechanism in the transformer architecture. our framework uses multi-head selfattention to decide which tokens and heads to prune in each layer. then, we review power-bert and length-adaptive transformer, which are recent works that eliminate tokens to improve efficiency in bert-based models. ## token and head adaptive transformer in this section, we describe the pruning strategy based on attention importance and head sensitivity score to remove tokens and attention heads respectively. we train token and head adaptive transformers with the pruning strategy to make the final models robust to arbitrary attention head and token drops at inference time. fig. 1 illustrates how tokens and heads are jointly pruned in transformer models. then we conduct a multi-objective evolutionary search on the trained model to find the optimal pruning configuration that meets the target computational budget. ## experiments we conduct extensive experiments to evaluate our framework and compare to the baseline models and length-adaptive transformers. ## results and analysis in this section, we evaluate the experiments and analyze the pruning results in terms of the accuracyefficiency trade-offs and head distribution. ## related work transformers and bert have outperformed existing language based models and achieved state-ofthe-art results but at the cost of high computational power requirements. in order to make these models efficient in terms of memory and power footprint, lat sequence tha sequence tha head lat sequence tha sequence tha head lat sequence tha sequence tha head lat sequence tha sequence tha head lat sequence tha sequence tha head lat sequence tha sequence tha head figure 3: comparison of the sequence length configurations of the ⋆ models with the maximum efficiency gains on squad and sst-2 benchmarks. lat sequence represents the token length configuration from length-adaptive transformer while tha for token and head adaptive transformers. tha head shows the retaining head distribution of token and head adaptive models. while large networks are essential for the best performance, the model size becomes an obstacle when it comes to scaling and also leads to worse training time. albert @xcite targets this limitation and proposed parameter reduction training by factoring the embedding matrices and sharing the parameters across layers which resulted in a smoother up-scaling up of the base model as well as significantly less training time. however, in terms of inference latency, there is not much improvement. power-bert @xcite points out this fact and proposes an orthogonal approach of progressive fine-grained word-vector elimination that can reduce the inference time up to 4.5x without harming the accuracy much. however, power-bert requires repetitive training for each different constraint and is limited to sequencelevel classification only. length adaptive transformer or lat @xcite alleviates these shortcomings by introducing a lengthdrop method that automatically derives multiple submodels using evolutionary search without requiring any re-training. they also proposed a drop-and-restore process to set aside the word-vectors so that they can be restored at the final layers extending the power-bert beyond classification tasks to a broader range of nlp tasks. parallel to these studies that deal with token pruning, another line of work investigates whether we need so many attention heads. @xcite identifies the most important attention heads using layer-wise relevance propagation or lrp @xcite and gradually prunes the head accordingly. in a complementary approach, @xcite finds the most relevant heads by masking or ablating one or more heads and looking at their impact on performance. in this way, they calculate a proxy importance score and iteratively prune the heads according to that without degrading the performance significantly. another group of studies concentrates on model architecture and parameter size. where distilbert @xcite uses a teacher-student setting to transfer knowledge from a large model to its small student version to reduce the model parameters, mobilebert @xcite reduces it by leveraging the bottleneck structure. while both token and head pruning improves the model's efficiency i.e. less inference time, lower memory, and power footprint, to the best of our knowledge, no study has tried to optimize both of them simultaneously which precisely lays down the basis of our work. ## conclusion and future work in this work, we propose a joint token and head pruning framework to train robust transformer models adaptive to token and head dropouts. our idea is to calculate head importance from the selfattention operation and eliminate redundant heads from the retained tokens to further reduce the computational cost and train the model adaptive and perform better generalization. the final models trained with the proposed framework show better generalization and greater efficiency gains. token and head adaptive transformers are robust to prune out more tokens than length-adaptive transformers along with the attention heads. this leads to superior accuracy-latency trade-offs and larger area-under-curve of the pareto frontier found with an evolutionary search. we would also like to explore further in the direction of reducing the computational cost of large-scale models by eliminating redundant operations and applying efficient optimization techniques to search for efficient models with minimum resources.