Blockwise Self-Attention for Long Document Understanding (2020)

Abstract: We present BlockBERT, a lightweight and efficient BERT model for better modeling long-distance dependencies. Our model extends BERT by introducing sparse block structures into the attention matrix to reduce both memory consumption and training/inference time, which also enables attention heads to capture either short- or long-range contextual information. We conduct experiments on language model pre-training and several benchmark question answering datasets with various 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 comparable and sometimes better prediction accuracy, compared to an advanced BERT-based model, RoBERTa.

## introduction recent emergence of the pre-training and finetuning paradigm, exemplified by methods like elmo @xcite , gpt-2/3 @xcite , bert @xcite , xlnet @xcite , roberta @xcite and albert @xcite , has drastically reshaped the landscape of the natural language processing research. these methods first pre-train a deep model with language model objectives 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-theart baselines across various tasks, such as question answering @xcite , coreference resolution @xcite , relation extraction @xcite ) and text retrieval @xcite , to name a few. ## background: memory bottleneck in training bert we briefly review bert and introduce its memory profiling in this section. following the paradigm of language model pre-training and down-stream task fine-tuning, bert @xcite consists of multiple layers of bidirectional transformers @xcite , where each transformer encoder has a multi-head self-attention layer and a position-wise feed-forward layer. using the same notation as in @xcite , we denote the number of transformer layers by l, the number of hidden units by h, the number of attention heads by a, the sequence length by n , and the batch size by b. we also assume the feed-forward hidden unit size to be 4h. foot_1 ## model: blockbert following @xcite , the dot-product attention in transformer is defined as: in this work, we design m to be a sparse block matrix, which not only reduces memory and the number of floating-point operations (flops) but also benefits from efficient dense matrix support from deep learning frameworks, such as pytorch and tensorflow. more formally, we split the length-n input sequence into n blocks, with each block of length n n . 4 the n ⇥ n attention matrix is then partitioned into n ⇥ n blocks, where each block matrix is of the size n n ⇥ n n . we define a sparse block matrix m by a permutation ⇡ of {1, 2, • • • , n}: in other words, blockbert reduces both o(n 2 ) memory consumption and flops by a factor of n, since n n ⇥ n n ⇥ @xmath0. ## experiments we evaluate the pre-training and fine-tuning performance of blockbert. in particular, when @xmath1, we denote 10:2 to be the configuration which assigns 10 heads to permutation (1, 2) and 2 to permutation (2, 1); when @xmath2, we denote 8:2:2 to be the configuration which assigns 8, 2, 2 heads to permutation (1, 2, 3), (2, 3, 1), and (3, 1, 2), respectively. we compare blockbert with the following baselines: google bert google bert is the official pretrained model from @xcite . ## related work in this section, we review the related work of memory optimization for neural network training and recent efforts to simplify transformer and bert. ## conclusion in this work, we study the lightweight bert model with the goal of achieving both efficiency and effectiveness. we profile and analyze the memory bottlenecks of bert and focus on optimize dotproduct self-attention, which consumes quadratic memory with respect to the sequence length. to reduce both time and memory consumption, we present blockbert, which sparsifies the attention matrices to be sparse block matrices. the proposed model achieves time and memory saving without significant loss of performance. in the future, we plan to benchmark more efficient transfomers in language model pre-training and fine-tuning. we also would like to explore more applications of blockbert on nlp tasks involving long sequences such as coreference resolution @xcite and document-level machine translation @xcite , and also non-nlp tasks such as protein sequence modeling @xcite .