Blockwise Self-Attention for Long Document Understanding

1 INTRODUCTION . Recent emergence of the pre-training and fine-tuning paradigm , exemplified by methods like ELMo ( Peters et al. , 2018 ) , GPT-2 ( Radford et al. , 2019 ) , BERT ( Devlin et al. , 2019 ) , XLNet ( Yang et al. , 2019 ) and RoBERTa ( Liu 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 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 reestablished the new state-ofthe-art baselines across various tasks , such as question answering ( Devlin et al. , 2019 ) , coreference resolution ( Joshi et al. , 2019b ) , relation extraction ( Soares et al. , 2019 ) and text retrieval ( Lee et al. , 2019 ; Nogueira & Cho , 2019 ) , to name a few . Building such models in practice , however , is an extremely resource-intensive process . For instance , the training of BERT-family models is notoriously expensive . Devlin et al . ( 2019 ) report that it takes four days for pre-training BERT-Base/BERT-Large on 4/16 Cloud TPUs , respectively . 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 that contributes 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 consumption , our analysis ( Section 2.1 ) shows that one of the main bottlenecks is actually dot-product selfattention , operated in multiple layers of Transformers ( Vaswani et al. , 2017 ) , the building block of BERT . As the attention operation is quadratic to the sequence length , this fundamentally limits the maximum length of the input sequence , and thus restricts the model capacity in terms of capturing long-distance dependencies . As a result , downstream tasks have to either truncate their sequences to leading tokens ( Nogueira & Cho , 2019 ) or split their sequences with a sliding window ( Joshi et al. , 2019a ; b ) . Ad-hoc handling of long sequences is also required in the pre-training stage , such as updating the model using only short sequences in the early stage ( Devlin et al. , 2019 ) . Common strategies for reducing memory consumption , unfortunately , do not work . For instance , shrinking the model by lowering the number of layers L , attention heads A , or hidden units H leads to significant performance degradation ( Vaswani et al. , 2017 ; Devlin et al. , 2019 ) and does not address the long sequence issue . Alternatively , general low-memory training techniques , such as 1https : //github.com/google-research/bert microbatching ( Huang et al. , 2018 ) and gradient checkpointing ( Chen et al. , 2016 ) essentially trade off training time for memory consumption , prolongs the already lengthy training process . In this work , we explore a different strategy , sparsifying the attention layers , intending to design a lightweight and effective BERT that can model long sequences in a memory-efficient way . Our BlockBERT extends BERT by introducing sparse block substructures into the attention matrix to reduce both memory consumption and the number of floating point operations ( FLOPs ) , which also enables attention heads to capture either short- or long-range contextual information . Compared to the previous method that also enforces sparsity ( e.g. , Child et al . ( 2019 ) ) , our approach is much simpler mathematically and very easy to implement . More importantly , the results of experiments conducted on several benchmark question answering datasets with various paragraph lengths show that BlockBERT performs comparably or even better than the original BERT-family models , while enjoying an 18.7-36.1 % reduction in memory usage and 12.0-25.1 % reduction in training time . The rest of the paper is organized as follows . Section 2 gives a brief introduction of the BERT model , along with an in-depth analysis of its memory usage during training time . We describe our proposed model in Section 3 and contrast it with existing methods that aim for creating a lighter model . Section 4 presents the experimental results and ablation studies , followed by a short survey of other related work in Section 5 and the conclusion in Section 6 . 2 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 ( Devlin et al. , 2019 ) consists of multiple layers of bidirectional Transformers ( Vaswani et al. , 2017 ) , where each Transformer encoder has a multi-head self-attention layer and a position-wise feed-forward layer . Using the same notation as in ( Devlin et al. , 2019 ) , 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 .2 2.1 MEMORY PROFILING . Training BERT is a memory-intensive process . In order to identify the bottleneck , we follow the memory model proposed by Sohoni et al . ( 2019 ) , where the memory usage throughout neural network training is categorized into three main types : ( 1 ) Model Memory is used to store model parameters ; ( 2 ) Optimizer Memory is the additional memory used by the specific learning algorithm during the process ; ( 3 ) Activation Memory consists of the outputs of each layer , which are cached for reuse in backpropagation to compute gradients . Take BERT-Base training as an example . The model has 110M parameters , so the model memory uses 0.2 GB if stored in half-precision floating-point format ( FP16 ) . For Adam ( Kingma & Ba , 2014 ) , the optimizer needs additional memory to store the gradients , first moments , and second moments of model parameters . If stored using the same precision , the optimizer memory should be three times of model memory.3 To calculate the exact size of activation memory is not trivial because it depends heavily on the implementation of the toolkit . Instead , we measure it empirically by training BERT-Base using Adam with a memory profiler ( more details are provided in Appendix A.2 ) . We use 32 NVIDIA V100 GPUs for training . Each single GPU thus consumes a mini-batch of size b = B/32 = 8 . Figure 1a shows the profiling result for a single GPU , where the model/optimizer/activation memory consumes 0.21/1.03/8.49 GB , resp . We can see that activation memory accounts for the vast majority of the total GPU memory ( 87.6 % ) and is clearly the bottleneck . Notice that although our analysis is done on BERT-Base , it can be easily generalized to BERT-Large and other models such as RoBERTa ( Liu et al. , 2019 ) and XLNet ( Yang et al. , 2019 ) . 2The default parameter settings for BERT-Base and BERT-Large can be found in Table 5 in Appendix A.1 . 3In the current PyTorch Adam implementation , the first and second moments are stored in single precision . Consequently , BERT ’ s optimizer memory ( 1 GB ) is five times of model memory ( 0.2 GB ) . 2.2 A REGRESSION ANALYSIS ON ACTIVATION MEMORY . For BERT , or more specifically , Transformer , the activation memory corresponds to intermediate results of different layers It grows linearly in all the model hyper-parameters , except the sequence length N , due to the attention layers . To quantify more clearly the O ( N ) and O ( N2 ) components in the activation memory , we conduct a regression analysis as follows . Assume that the activation memory ( in each GPU ) is a polynomial a2bN2 + a1bN + a0 , where b is the batch size in each GPU . If we fix the total number of tokens in a GPU , i.e. , b×N , to be constant ( in our case , 4096 ) , we should have a linear function w.r.t . N , i.e. , 4096a2N + 4096a1 + a0 . We enumerate N from { 128 , 256 , 512 , 1024 } in our experiments , and plot the corresponding profiled activation memory in Figure 1b . Using ordinary least squares ( OLS ) , with b × N = 4096 , the estimated linear function for activation memory is 0.00715 × N + 4.83 , where the first term is responsible for the O ( N2 ) component . When N = 512 , we can see that for BERT-Base , the O ( N2 ) component accounts for 3.66 GB and O ( N ) accounts for 4.83 GB . When the sequence length N increases to 1024 , however , the O ( N2 ) component increases to 7.32 GB , while O ( N ) is unchanged . 2.3 GENERAL TECHNIQUES FOR REDUCING MEMORY USAGE IN MODEL TRAINING . Observing that activation memory is the bottleneck , we discuss the effectiveness of common memory reduction techniques for BERT training below . Low Precision ( Micikevicius et al. , 2017 ) : Low precision is to use half-precision or mixed-precision for training neural networks . This technique has been widely used in Transformer training ( Ott et al. , 2019 ; Liu et al. , 2019 ) . In this work , we already assume to use mixed-precision training by default , as indicated in the aforementioned analysis . Microbatching ( Huang et al. , 2018 ) : Microbatching is to split a batch into small microbatches ( which can be fit into memory ) , and then run forward and backward passes on them separately with gradients for each micro-batch accumulated . Because it runs forward/backward pass multiple times for a single batch , it trades off time for memory . Gradient Checkpointing ( Chen et al. , 2016 ) : Gradient checkpointing saves memory by only caching activations of a subset of layers . The un-cached activations will be recomputed during backpropagation from the latest checkpoint . This strategy trades off time for memory by repeating computations that require large memory and will obviously extend the model training time . Knowledge Distillation ( Hinton et al. , 2015 ) : Knowledge distillation aims to compress and transfer knowledge from a teacher model to a simpler student model . However , knowledge distillation relies on a teacher model ( which is still expensive in training time ) and usually suffers from a certain degree of performance degradation . As common techniques are limited in reducing both the training time and memory usage , we investigate how to optimize the dot-product attention layers and introduce our approach next . 3 MODEL : BLOCKBERT . Following ( Vaswani et al. , 2017 ) , the dot-product attention in Transformer is defined as : Attention ( Q , K , V ) = softmax ( QK > √ d ) V . where Q , K , V ∈ RN×d with N to be the sequence length and d to be a hidden dimension . As we can see , the inner product between Q and K consumes O ( N2 ) memory . One simple way to reduce memory consumption of attention is to sparsify the attention matrix . Suppose we have a masking matrix M ∈ { 0 , 1 } N×N . We define a masked version of attention as follows : Attention ( Q , K , V , M ) = softmax ( QK > √ d M ) V , ( 1 ) with operator defined by ( A M ) ij = { Aij if Mij = 1 −∞ if Mij = 0 . 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 partitions , with each partition of length Nn . 4 The N × N attention matrix is then partitioned into n × n blocks , where each block matrix is of size Nn × N n . A sparse block matrix M can be defined by a permutation π of { 1 , 2 , · · · , n } : Mij = { 1 if π ( b ( i−1 ) n N + 1c ) = b ( j−1 ) n N + 1c 0 otherwise . ( 2 ) By writing Q , K , V as be block matrices , such that Q = [ Q > 1 · · · Q > n ] > , K = [ K > 1 · · · K > n ] > and V = [ V > 1 · · · V > n ] > and pluging them into Equation 1 , we can formally define Blockwise Attention as follows : Blockwise-Attention ( Q , K , V , M ) = softmax ( Q1K > π ( 1 ) √ d ) Vπ ( 1 ) ... softmax ( QnK > π ( n ) √ d ) Vπ ( n ) . ( 3 ) As a result , it only needs to compute and store QiK > π ( i ) ( i = 1 , · · ·n ) , each of which has size N n × N n . In other words , BlockBERT reduces the corresponding O ( N 2 ) memory consumption and FLOPs by a factor of n , since Nn × N n × n = N×N n .

Summary: The paper propose to sparsify the attention matrix to decrease memory usage and to speed up training. The authors experiment the model on multiple tasks. The model gains ~20% efficiency with ~20% decrease in memory use while maintaining comparable performance to the state of the art model. To keep the performance comparable, the authors use the same training corpus. The authors also discuss how block size could change the performance of the model. The paper is clear and well organized with good experiment results.