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

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.

Efficient Attention via Pruned Token Compression for Document Classification

The paper introduces an innovative approach to address the computational inefficiencies of transformer-based models like BERT, which are prevalent in natural language processing tasks but are hindered by their computationally expensive self-attention mechanisms. The authors propose a method that integrates token pruning and token combining strategies to enhance both performance and efficiency in document classification tasks.

Key Contributions

1. Token Pruning: The paper presents a novel token pruning strategy that selectively removes less important tokens from the attention mechanism’s key and value matrices. This approach is augmented by implementing fuzzy logic, which helps to manage uncertainty and mitigate risks associated with potential mispruning due to the imbalanced distribution of token importance.
2. Token Combining: Complementing the pruning strategy, the authors propose a token combining method that reduces the size of input sequences by condensing tokens. Employing ideas from Slot Attention, this method further compresses the model, enhancing efficiency without significant information loss.
3. Experimental Validation: Empirical evidence demonstrates the effectiveness of these integrated strategies across various datasets, achieving notable improvements over the standard BERT model with an increase of 5% in accuracy and 5.6% in F1 score, alongside a substantial reduction in memory cost (0.61x) and speedup (1.64x).

Implications and Future Work

The integration of pruning and combining not only addresses the computational bottleneck inherent in BERT and similar models but also provides a template for enhancing the efficiency of other transformer architectures. By focusing on the core useful tokens, this work suggests pathways for more extensive applications in long-document processing and other NLP tasks requiring resource efficiency.

Future research could build upon this work by exploring the adaptability of the proposed framework to other transformer variants and expanding it to include additional NLP tasks beyond document classification. Moreover, the introduction of fuzzy logic in neural architecture design opens up avenues for its application in other areas of machine learning where uncertainty is a concern.

Conclusion

The paper makes significant strides in improving the efficiency and effectiveness of transformer models in document classification. By leveraging token pruning and combining, the proposed method reduces computation costs while maintaining or exceeding the performance of existing models. This dual-strategy approach may serve as a foundation for future innovations aimed at optimizing neural network models for efficient real-world applications.