Fine- and Coarse-Granularity Hybrid Self-Attention for Efficient BERT (2022)

Abstract: Transformer-based pre-trained models, such as BERT, have shown extraordinary success in achieving state-of-the-art results in many natural language processing applications. However, deploying these models can be prohibitively costly, as the standard self-attention mechanism of the Transformer suffers from quadratic computational cost in the input sequence length. To confront this, we propose FCA, a fine- and coarse-granularity hybrid self-attention that reduces the computation cost through progressively shortening the computational sequence length in self-attention. Specifically, FCA conducts an attention-based scoring strategy to determine the informativeness of tokens at each layer. Then, the informative tokens serve as the fine-granularity computing units in self-attention and the uninformative tokens are replaced with one or several clusters as the coarse-granularity computing units in self-attention. Experiments on the standard GLUE benchmark show that BERT with FCA achieves 2x reduction in FLOPs over original BERT with <1% loss in accuracy. We show that FCA offers a significantly better trade-off between accuracy and FLOPs compared to prior methods.

## introduction transformer-based large pre-trained language models with bert @xcite as a typical model routinely achieve state-of-the-art results on a number of natural language processing tasks @xcite @xcite , such as sentence classification @xcite , question answering @xcite , and information extraction @xcite . despite notable gains in accuracy, the high computational cost of these large models slows down their inference speed, which severely impairs their practicality, especially in the case of limited industry time and resources, such as mobile phone and aiot. in addition, the excessive energy consumption and environmental impact caused by the computation of these models also raise the widespread concern @xcite . to improve the efficiency of bert, the mainstream techniques are knowledge distillation @xcite and pruning. knowledge distillation aims to transfer the "knowledge" from a large teacher model to a lightweight student model. the student model is then used during inference, such as distilbert @xcite . pruning technique includes: (1) structured methods that prune structured blocks of weights or even complete architectural components in bert, for example encoder layers @xcite , (2) unstructured methods that dynamically drop redundant units, for example, attention head @xcite and attention tokens @xcite . however, both types of methods encounter challenges. for the former, a great distillation effect often requires an additional large teacher model and very complicated training steps @xcite . for the latter, pruning methods discard some computing units, which inevitably causes information loss. in contrast to the prior approaches, we propose a self-motivated and information-retained technique, namely fca, a fine-and coarse-granularity hybrid self-attention that reduces the cost of bert through progressively shortening the computational sequence length in self-attention. specifically, fca first evolves an attention-based scoring strategy to assign each token with the informativeness. through analyzing the informativeness distribution at each layer, we conclude that maintaining fulllength token-level representations is progressive redundant along with layers, especially for the classification tasks that only require single-vector repre-sentations of sequences. consequently, the tokens are divided into informative tokens and uninformative tokens according to their informativeness. then, they are updated through different computation paths. the informative tokens carry most of the learned features and remain unchanged as the fine-grained computing units in self-attention. the uninformative tokens may not be as important as informative ones but we will not completely discard them to avoid information loss. instead, we replace them with more efficient computing units to save memory consumption. experiments on the standard glue benchmark show that fca accelerates bert inference speed and maintains high accuracy as well. our contributions are summarized as follows: ## related work there has been much prior literature on improving the efficiency of transformers. the most common technologies include: knowledge distillation refers to training a smaller student model using outputs from various intermediate representations of larger pre-trained teacher models. in the bert model, there are multiple representations that the student can learn from, such as the logits in the final layer, the outputs of the encoder units, and the attention maps. the distillation on output logits is most commonly used to train smaller bert models @xcite @xcite . the output tensors of encoder units contain meaningful semantic and contextual relationships between input tokens. some work creates a smaller model by learning from the outputs of teacher's encoder @xcite @xcite . attention map refers to the softmax distribution output of the self-attention layers and indicates the contextual dependence between the input tokens. a common practice of distillation on attention maps is to directly minimize the difference between the self-attention outputs of the teacher and the student @xcite @xcite . this line of work is orthogonal to our approach and our proposed fca can be applied to the distillate models to further accelerate their inference speed. pruning refers to identifying and removing less important weights or computation units. pruning methods for bert broadly fall into two categories. unstructured pruning methods prune individual weights by comparing their absolute values or gradients with a pre-defined threshold @xcite @xcite . the weights lower than the threshold are set to zero. unlike unstructured pruning, structured pruning aims to prune structured blocks of weights or even complete architectural components in the bert model. @xcite pruned attention heads using a method based on stochastic gates and a differentiable relaxation of the l0 penalty. @xcite randomly dropped transformer layers to sample small sub-networks from the larger model during training which are selected as the inference models. @xcite progressively reduced sequence length by pruning word-vectors based on the attention values. this work is partly similar to the fine-grained computing units in our proposed fca. however they ignored the coarse-grained units that may cause information loss. in addition, there are some engineering techniques to speed up the inference speed, such as mixed precision @xcite and quantization @xcite . using half-precision or mixed-precision representations of floating points is popular in deep learning to accelerate training and inference speed. quantization refers to reducing the number of unique values required to represent the model weights, which in turn allows to represent them using fewer bits. ## preliminary bert @xcite is a transformerbased language representation model, which can be fine-tuned for many downstream nlp tasks, including sequence-level and token-level classification. the transformer architecture @xcite is a highly modularized neural network, where each transformer layer consists of two submodules, namely the multi-head self-attention sublayer (mha) and the position-wise feed-forward network sub-layer (ffn). both sub-modules are wrapped by a residual connection and layer normalization. mha. the self-attention mechanism allows the model to identify complex dependencies between the elements of each input sequence. it can be formulated as querying a dictionary with key-value pairs. formally, previous research @xcite has shown that in addition to mha sub-layer, ffn sub-layer also consumes large memory in terms of model size and flops. as a result, if we reduce the computational sequence length of mha, the input and the consumption of ffn sub-layer will become less accordingly. ## methodologies to shorten the computational sequence length of self-attention, our core motivation is to divide tokens into informative and uninformative ones and replace the uninformative tokens with more efficient units. this section introduces each module of our model in detail. ## discussion our proposed fca is dedicated to the classification tasks that only require single-vector representations, and it can not be directly applied to the tasks of requiring to maintain the full input sequence in the output layer, such as ner and extractive mrc. on these tasks, we need to make some modifications of only performing fca operation over k and v in self-attention and maintaining the full length of q. the eq. 2 is modified to: we also attempted to maintain the full length of k and v and shorten q, but the experimental results are unsatisfactory. ## conclusion in this paper, we propose fca, a fine-and coarsegranularity hybrid self-attention that reduces the computation cost through progressively shortening the computational sequence length in self-attention. experiments on glue and race datasets show that bert with fca achieves 2x reduction in flops over original bert @xmath0. meanwhile, fca offers significantly better trade-off between accuracy and flops compared to prior methods.