Focus on the Core: Efficient Attention via Pruned Token Compression for Document Classification (2023)

Abstract: Transformer-based models have achieved dominant performance in numerous NLP tasks. Despite their remarkable successes, pre-trained transformers such as BERT suffer from a computationally expensive self-attention mechanism that interacts with all tokens, including the ones unfavorable to classification performance. To overcome these challenges, we propose integrating two strategies: token pruning and token combining. Token pruning eliminates less important tokens in the attention mechanism’s key and value as they pass through the layers. Additionally, we adopt fuzzy logic to handle uncertainty and alleviate potential mispruning risks arising from an imbalanced distribution of each token’s importance. Token combining, on the other hand, condenses input sequences into smaller sizes in order to further compress the model. By integrating these two approaches, we not only improve the model’s performance but also reduce its computational demands. Experiments with various datasets demonstrate superior performance compared to baseline models, especially with the best improvement over the existing BERT model, achieving +5%p in accuracy and +5.6%p in F1 score. Additionally, memory cost is reduced to 0.61x, and a speedup of 1.64x is achieved.

## introduction transformer-based deep learning architectures have achieved dominant performance in numerous areas of natural language processing (nlp) studies @xcite @xcite . in particular, pre-trained transformer-based language models like bert @xcite and its variants @xcite @xcite have demonstrated state-of-the-art performance on many nlp tasks. the self-attention mechanism, a key element in transformers, allows for interactions between every pair of tokens in a sequence. this effectively captures contextual information across the entire sequence. this mechanism has proven to be particularly beneficial for text classification tasks @xcite @xcite . despite their effectiveness, bert and similar models still face major challenges. bert can be destructive in that not all tokens contribute to the final classification prediction @xcite . not all tokens are attentive in multi-head self-attention, and uninformative or semantically meaningless parts of the input may not have a positive impact on the prediction @xcite . further, the self-attention mechanism, which involves interaction among all tokens, suffers from substantial computational costs. its quadratic complexity relative to the length of the input sequences results in high time and memory costs, making training impractical, especially for document classifications @xcite . in response to these challenges, many recent studies have attempted to address the problem of computational inefficiency and improve model ability by focusing on a few core tokens, thereby reducing the number of tokens that need to be processed. their intuition is similar to human reading comprehension achieved by paying closer attention to important and interesting words @xcite . one approach is a pruning method that removes a redundant token. studies have shown an acceptable trade-off between performance and cost by simply removing tokens from the entire sequence to reduce computational demands @xcite @xcite . however, this method causes information loss, which degrades the performance of the model @xcite . unlike previous studies, we apply pruning to remove tokens from the keys and values of the attention mechanism to prevent the information loss and reduce the cost. in our method, less important tokens are removed, and the number of tokens gradually decreases by a certain ratio as they pass through the layers. however, there is still a risk of mispruning when the distribution of importance is imbalanced, since the ratio-based pruning does not take into account the importance distribution @xcite . to address this issue, we propose to adopt the fuzzy logic by utilizing fuzzy membership functions to reflect the uncertainty and support token pruning. however, the trade-off between performance and cost of pruning limits the number of tokens that can be removed, hence, self-attention operations may still require substantial time and memory resources. for further model compression, we propose a token combining approach. another line of prior works @xcite @xcite have demonstrated that combining tokens can reduce computational costs and improve performance in various computer vision tasks, including image classification, object detection, and segmentation. motivated by these studies, we aim to compress text sequence tokens. since text differs from images with locality, we explore slot attention @xcite , which can bind any object in the input. instead of discarding tokens from the input sequence, we combine input sequences into smaller number of tokens adapting the slot attention mechanism. by doing so, we can decrease the amount of memory and time required for training, while also minimizing the loss of information. in this work, we propose to integrate token pruning and token combining to reduce the computational cost while improving document classification capabilities. during the token pruning stage, less significant tokens are gradually eliminated as they pass through the layers. we implement pruning to reduce the size of the key and value of attention. subsequently, in the token combining stage, tokens are merged into a combined token. this process results in increased compression and enhanced computational efficiency. we conduct experiments with document classification datasets in various domains, employing efficient transformer-based baseline models. compared to the existing bert model, the most significant improvements show an increase of 5%p in accuracy and an improvement of 5.6%p in the f1 score. additionally, memory cost is reduced to 0.61x, and a speedup of 1.64x is achieved, thus accelerating the training speed. we demonstrate that our integration results in a synergistic effect not only improving performance, but also reducing memory usage and time costs. our main contributions are as follows: ## methods in this section, we first describe the overall architecture of our proposed model, which integrates token pruning and token combining. then, we introduce each compression stage in detail, including the token pruning strategy in section 3.2 and the token combining module in section 3.3. ## experiments this section aims to validate the effectiveness of our proposed model. firstly, we evaluate the document classification performance of our proposed model compared to the baseline models. secondly, we investigate the time and memory costs of our proposed model and evaluate its efficiency. lastly, through the ablation study, we compare the effects of different preservation ratios on fuzzy-based token pruning self-attention. we also analyze the impact of the position of the token combining module and the number of combination tokens. ## conclusion in this paper, we introduce an approach that integrates token pruning and token combining to improve document classification by addressing expensive and destructive problems of self-attention in the existing bert model. our approach consists of fuzzy-based token pruned attention and token combining module. our pruning strategy gradually removes unimportant tokens from the key and value in attention. moreover, we enhance the robustness of the model by incorporating fuzzy membership functions. for further compression, our token combining module reduces the time and memory costs of the model by merging the tokens in the input sequence into a smaller number of combination tokens. experimental results show that our proposed model enhances the document classification performance by reducing computational requirements with focusing on more significant tokens. our findings also demonstrate a synergistic effect by integrating token pruning and token combining, commonly used in object detection and semantic segmentation. ultimately, our research provides a novel way to use pre-trained transformer models more flexibly and effectively, boosting performance and efficiency in a myriad of applications that involve text data processing.