Blockwise Self-Attention for Long Document Understanding

Findings of the Association for Computational Linguistics: EMNLP 2020, pages 2555–2565 November 16 - 20, 2020. c⃝2020 Association for Computational Linguistics 2555 Blockwise Self-Attention for Long Document Understanding Jiezhong Qiu1⇤, Hao Ma2, Omer Levy2, Wen-tau Yih2, Sinong Wang2, Jie Tang1 1Department of Computer Science and Technology, Tsinghua University 2Facebook AI qiujz16@mails.tsinghua.edu.cn {haom,omerlevy,scottyih,sinongwang}@fb.com jietang@tsinghua.edu.cn Abstract We present BlockBERT, a lightweight and ef- ficient BERT model for better modeling long- distance dependencies. Our model extends BERT by introducing sparse block structures into the attention matrix to reduce both mem- ory consumption and training/inference time, which also enables attention heads to cap- ture either short- or long-range contextual in- formation. We conduct experiments on lan- guage model pre-training and several bench- mark question answering datasets with vari- ous paragraph lengths. BlockBERT uses 18.7- 36.1% less memory and 12.0-25.1% less time to learn the model. During testing, BlockBERT saves 27.8% inference time, while having com- parable and sometimes better prediction accu- racy, compared to an advanced BERT-based model, RoBERTa. 1 Introduction Recent emergence of the pre-training and fine- tuning paradigm, exemplified by methods like ELMo (Peters et al., 2018), GPT-2/3 (Radford et al., 2019; Brown et al., 2020), BERT (Devlin et al., 2019), XLNet (Yang et al., 2019), RoBERTa (Liu et al., 2019) and ALBERT (Lan et al., 2019), has drastically reshaped the landscape of the natural language processing research. These methods first pre-train a deep model with language model objec- tives using a large corpus and then fine-tune the model using in-domain supervised data for target applications. Despite its conceptual simplicity, this paradigm has re-established the new state-of-the- art baselines across various tasks, such as question answering (Devlin et al., 2019), coreference resolu- tion (Joshi et al., 2019b), relation extraction (Soares et al., 2019) and text retrieval (Lee et al., 2019; Nogueira and Cho, 2019), to name a few. ⇤This work was partially done when the first author was an intern at Facebook AI. Code is available at https:// github.com/xptree/BlockBERT Building such models in practice, however, is an extremely resource-intensive process. For in- stance, the training of BERT-family models is noto- riously expensive. Devlin et al. (2019) report that it takes four days to pre-train BERT-Base/BERT- Large on 4/16 Cloud TPUs. In order to reduce the pre-training time of RoBERTa to 1 day, Liu et al. (2019) use 1,024 V100 GPUs. One crucial factor contributing to the long training time is the memory consumption of these deep models, as it directly affects the batch size. Although the fine-tuning stage is relatively inexpensive, the memory issue still restricts the scenarios in which BERT can be used. For instance, “it is currently not possible to re-produce most of the BERT-Large results on the paper using a GPU with 12GB-16GB of RAM, because the maximum batch size that can fit in memory is too small.1” Although one may think that model size is the main contributor to the large memory consump- tion, our analysis (Section 2.1) shows that one of the main bottlenecks is actually dot-product …

Content truncated to 500 words