Fine- and Coarse-Granularity Hybrid Self-Attention for Efficient BERT

Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 4811 - 4820 May 22-27, 2022 c⃝2022 Association for Computational Linguistics Fine- and Coarse-Granularity Hybrid Self-Attention for Efficient BERT Jing Zhao, Yifan Wang, Junwei Bao∗, Youzheng Wu, Xiaodong He JD AI Research, Beijing, China {zhaojing857,wangyifan15,baojunwei, wuyouzheng1,xiaodong.he}@jd.com Abstract Transformer-based pre-trained models, such as BERT, have shown extraordinary success in achieving state-of-the-art results in many natu- ral language processing applications. However, deploying these models can be prohibitively costly, as the standard self-attention mechanism of the Transformer suffers from quadratic com- putational cost in the input sequence length. To confront this, we propose FCA, a fine- and coarse-granularity hybrid self-attention that re- duces the computation cost through progres- sively shortening the computational sequence length in self-attention. Specifically, FCA con- ducts an attention-based scoring strategy to de- termine 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 re- placed with one or several clusters as the coarse- granularity computing units in self-attention. Experiments on GLUE and RACE datasets show that BERT with FCA achieves 2x reduc- tion in FLOPs over original BERT with <1% loss in accuracy. We show that FCA offers significantly better trade-off between accuracy and FLOPs compared to prior methods1. 1 Introduction Transformer-based large pre-trained language mod- els with BERT (Devlin et al., 2018) as a typical model routinely achieve state-of-the-art results on a number of natural language processing tasks (Yang et al., 2019; Liu et al., 2019; Clark et al., 2020), such as sentence classification (Wang et al., 2018), question answering (Rajpurkar et al., 2016, 2018), and information extraction (Li et al., 2020b). Despite notable gains in accuracy, the high com- putational cost of these large models slows down their inference speed, which severely impairs their ∗Corresponding author 1Code is available at https://github.com/ pierre-zhao/FCA-BERT practicality, especially in the case of limited indus- try time and resources, such as Mobile Phone and AIoT. In addition, the excessive energy consump- tion and environmental impact caused by the com- putation of these models also raise the widespread concern (Strubell et al., 2019; Schwartz et al., 2020). To improve the efficiency of BERT, the main- stream techniques are knowledge distillation (Hin- ton et al., 2015) and pruning. Knowledge distil- lation 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 (Sanh et al., 2019). Pruning technique includes: (1) structured methods that prune structured blocks of weights or even com- plete architectural components in BERT, for ex- ample encoder layers (Zhang and He, 2020), (2) unstructured methods that dynamically drop redun- dant units, for example, attention head (Voita et al., 2019) and attention tokens (Goyal et al., 2020). However, both types of methods encounter chal- lenges. For the former, a great distillation effect often requires an additional large teacher model and very complicated training steps (Jiao …

Content truncated to 500 words