Blockwise Self-Attention for Long Document Understanding
Authors: Jiezhong Qiu
Hao Ma
Omer Levy
Scott Wen-tau Yih
Sinong Wang
Jie Tang
Venue:  (2020)
Keywords: BERT
Transformer

Abstract: We present BlockBERT, a lightweight and efficient BERT model that is designed to 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 time, which also enables attention heads to capture either short- or long-range contextual information. We conduct experiments on several benchmark question answering datasets with various paragraph lengths. Results show that BlockBERT uses 18.7-36.1% less memory and reduces the training time by 12.0-25.1%, while having comparable and sometimes better prediction accuracy, compared to an advanced BERT-based model, RoBERTa.

Under review as a conference paper at ICLR 2020
BLOCKWISE SELF-ATTENTION FOR LONG DOCUMENT UNDERSTANDING
Anonymous authors Paper under double-blind review
ABSTRACT
We present BlockBERT, a lightweight and efficient BERT model that is designed to better modeling long-distance dependencies. Our model extends BERT by in- troducing sparse block structures into the attention matrix to reduce both memory consumption and training time, which also enables attention heads to capture ei- ther short- or long-range contextual information. We conduct experiments on sev- eral benchmark question answering datasets with various paragraph lengths. Re- sults show that BlockBERT uses 18.7-36. 1.36.% 1% less memory and reduces the training time by 12.0-25.11 %6, while having comparable and sometimes better prediction accuracy, compared to an advanced BERT-based model, RoBERTa.
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 al. 2019) and ROBERTa (Liu et al. 2019), has drastically reshaped the landscape of the natural language processing research. These methods first preitrain a deep model with language model ob- jectives using 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-of- the-art baselines across various tasks, such as question answering (Devlin 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 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- pre-training BERT- BERT-Baac/BERT-Large reduce the pre-training time of RoBERTa to day, Liu et al. (2019) use 1,024 V100 GPUs. One crucial factor that contributes 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 cur- rently not possible to re-produce most of the BERT-Large results on the paper using GPU with 12GB-16GB of RAM, because the maximum batch size that can fit in memory
4/16 Cloud TPUs, respectively. In order to
the long training time is the memory consumption of these deep
small f
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 self- attention, operated in multiple layers of Transformers (Vaswani et al. (0077). the building block of BERT. As the attention operation 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 result, downstream tasks have to either truncate their sequences to leading tokens (Nogueira & Cho. 2019) split their sequences with sliding window (Joshi et al. 20. b). Ad-hoc handling of long sequences is also required in the pre- mrairiinn stage, such updating the model using only short sequences in the early stage (Devlin et al. 2019).
Common strategics for reducing memory consumption, unfortunately, do not work. For instance, shrinking the model by lowering the number of layers L, attention heads A. hidden units H leads to significant performance degradation (Vaswani et al., 2017 Devlin et al. (019) and does not address the long sequence issue. Alternatively, general low-memory training techniques, such as

---

Under review as a conference paper at ICLR 2020
microbatching (Huang al., 2018) and gradient checkpointing (Chen al. 2016) essentially trade off training time for memory consumption, prolongs the already lengthy training process.
In this work, explore different strategy, sparsifying the attention layers, intending to design a lightweight and effective BERT that can model long sequences in a memory-efficiet 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 capture either short- or long-range contextual information. Compared to the previous method that also enforces sparsity (e.g.. Ch.ld et al. (2019)) our approach is much simpler mathematically and very easy to implement. More importandl, the results of experiments conducted on several benchmark question answering datasets with various paragraph lengths show that BlockBERR performs comparably or even better than the original BERT-family models, while enjoying an 18.7 187-336.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 brief introduction of the BERT model, along with in-depth analysis of its memory usage during training time. We describe our proposed model in Section 3 and contrast with existing methods that aim for creating a lighter model. Section [ presents the experimental results and ablation studies, followed by a short survey of other related work in Section and the conclusion in Section 6
2 BACKGROUND: MEMORY BOTTLENECK 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) con- sists of multiple layers of bidirectional Transformers (Vaswani et al. (017). where each Transformer encoder has a multi-head self-attention layer and a position-wise foed-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 4HH
2.1 MEMORY PROFILING
Training BERT is memory-intensive process. In order identify the bottleneck, we follow the memory model proposed by [Sohoni et al. (2019). where the memory usage throughout neural net- work training categorized into three main lypes: (I) Model Memory used store model parameters; (2) Optimizer Memory is the additional memory used by the specific learning algo- rithm 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- floating-poin format (FP16). For Adam (Kingma & Ba 2014]. the optimizer needs additional memory 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 train- ing 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 mini-batch of size b = B/32 = Figure Ia] shows the profiling result for a single GPU, where the memory consumes 0,211/.03/8.44 GB, resp. We can that activa- tion 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)
mode//optimizee/actiiaati
"The default parameter settings for BERT- Base and BERT-Large can be found in Table in Appendix A.1 II the current PyToch Adam implementation, the first and second moments are stored single precision. Consequently, BERT's optimizer memory (1 GB) five times of model memory (0.2 GB).
2

---

Under review as a conference paper at ICLR 2020
BERT-Base (GB)
BERT 0.00V15N
Act. Merorr (GB)
21% oytiminer 10 276
200
691
90
2005
87.6%
Sequence Length N
(a) BERT- BERT-Base Training Memory Profiling
(b) Regression Analysis on Activation Memory
Figure 1: Memory Profiling for BERT
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- nyper-parameters, except the sequence length N, due to the attention layers. To quantify more clearly the O(N) and O(N²) components in the activation memory, we conduct a regression analysis as follows. Assume that the activation memory (in each GPU) is polynomial azbN² + a bN + ag. where is the batch size in each GPU. we fix the total number of tokens a GPU, i.e., x N, to constant (in our case, 4096), we should have linear function w.r.t. N, i.e., 4096a, N + 4096a - og. We enumerate N from {128, 256, 512, 1024] in our experiments, and plot the corresponding profiled activation memory in Figure lb Using ordinary least squares (OLS), with b x N 4096, the estimated linear function for activation memory is 0.00715 x N + 4.83, where the first term is responsible for the O(N2) component. When N 512, 3.66 GB and ((N) accounts for 4.83 GB. When the sequence length N increases to 1024, however, the (NN component increases 7.32 GB, while O(N) is unchanged.
see that for BERT-Base the ONN) component accounts for
2.3 GENERAL TECHNIQUES FOR REDUCING MEMORY USAGE MODEL TRAINING
Observing that activation memory is the bottleneck, we discuss the effectiveness of common mem- ory reduction techniques for BERT training below.
Low Precision (Micikevicius et al. 2017): Low precision is to use half-precision mixed-precision for training neural networks. This technique has been widely used in Transformer training (Ou et al., 2019 Liu et al. 2019). this work, we already assume to use mixed -reed--reeiiion training by default, as indicated in the aforementioned analysis.
Microbatching (Huang et al.) 2018): Microbatching is split batch into small micro- batches (which can be fit into memory), and then run forward and backward passes on them sep- arately 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 subset of layers. The un-cached activations will be recomputed during buckpropagation 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 aims to compress and transfer knowledge from a teacher model 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.
Knowledge Distillation (Hinton et al. 2015):
As common techniques are limited in reducing both the training time and memory usage, we inves- tigate how to optimize the dot-product attention layers and introduce approach next.
3

---

Under review as a conference paper at ICLR 2020
3
MODEL: BLOCKBERT
Following Vassaan al. 2017). the dot-product attention in Transformer is defined as:
QK
Attention( Q, K. V) softmax
V.
Vd
where Q. K, V E Rxxd with N to be the sequence length and d to be a hidden dimension. As we can see, the inner product between Q and 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, NxN, We define masked version of attention follows:
QK
vd @ (M) GM)V.
Attention Q,K, MM softmax
(1)
with operator @ defined by
Ai
if M =1 if M. =
(A M);
-00
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 split the length-N input sequence into partitions, with each partition of length A The N N attention matrix is then partitioned into n n blocks, where each block matrix is of size N. A sparse block matrix M can be defined by permutation of {1, ... ,n]:
from deep learning frameworks, such as PyTorch and Tensorflow. More formally,
m
Mj 1 if (1 (1"-س\- + 11) (f-1) 12m + +1]
(2)
otherwise.
By writing Q, K, V as be block matrices, such that Q ToT and [v] v.]] and pluging them into Equation Attention as follows:
K = kT
we can formally define Blockwise
Q. Q va
softmax
V=(1)
Bleckwise-Attention(.. (, K,V,
(3)
9n.K min)
softmax
(n(
/dd
As a result, it only needs to compute and store Q.K #()) (i = 1,... each of which has size N x N. In other words, BlockBERT reduces the corresponding O(N?) memory consumption and n
FLOPs by a factor of n, since N N X =n= NxN
3.1 BLOCKWISS MULTI-HEAD ATTENTION
Analogous Multi-head Attention (Vawwani al. 2017). allow queries, keys, and values be projected multiple times and perform blockwise attentions in parallel. Moreover, different block- wise attention heads can use different masking matrices. The outputs of multiple heads are then concatenated and aggregated with another linear projection. Let A be the number of attention heads and H the number of hidden units. Blockwise multi-head attention is formally defined by:
Blockwise-Multi-heuee Attention Q, K, V) Concat(head,, )
where head, Blockwise-Attention QW? KW. VW!. M.),
with d A W W.K. W. uxd and the projection matrix W H Each masking
matrix M, is determined by permutation म; according to Equation 2 In particular, we choose from permutations generated by shifting one position:
(2,3, ................................ n, 1), i.e., we select For example, with 12 attention heads (A 12) and 2 blocks (n 2), one
{a, (a,
"We assume N can be divided by not, we pad the input sequence make N divisible.

---

Under review as a conference paper at ICLR 2020
Linear
Concal
n=2
Blockwise Attention
(1.2)
12.1)
Linear
Linear
Lincar
01.
18
Masking Matrices
Q
K
V
Figure 2: Architecture of Blockwise Multi-head Attention.
DERT DUKTIAX
BINTEER n.d BO0ET 4.M4
BlcSBERT and EANOBEY +.8
4.5
Mem. (GB)
N
Model
O(N)
O(N²)
Memory (GB)
BERT
4.83
3,66
512
8
BlockBERT n=2
4.84
1.83
BlockBERT n=3
4.87
1.22
BERT
4.83
7.32
1024
4 BlockBERT n=2
4.84
3.66
BlockBERT n=3
4.87
2.44
10.0
Sequence Length
N
Figure Regression analysis on activation memory for BERT and BlockBERT.
Table I: Estimated O(N2) and O(N) activation memory for BERT and BlockBERT.
configuration can be assigning 10 heads to permutation (1,2) and the other heads to permutation
(2, 1). Figure 2 illustrates the blockwise multi-head attention with the block numbers n (2,3). Blockwise sparsity captures both local and long-distance dependencies in memory-efficiency way, which crucial for long-docn mong-document understanding tasks. For instance, the identity permutation, ie., 21, n), enables each token to attend its nearby tokens in self-attention. Tokens within the same local group attend a long-distance group of tokens together in other permutations. Our proposed BRO essentially replaces the multi-head attention layers in TrRT with blockwise multi-head attention.
3.2 ANALYSIS OF MEMORY USAGE REDUCTION
To validate our claim that BlockBERT ERT with n x n blocks can reduce the O(N2) memory use by a fac-
tor of n, perform the same memory profiling described in sections 2.J and 2.2 Again, We fix the number of tokens in each GPU (bxN 4096) and choose N from { {128, 256 512. 1024, 2048 As
from Figure 3 and Table I the empirical results align well with the theoretical values. When we set block size to be 2 and for BlockBERT, their estimated ( activation memory decreases to 1/2 and 1/3 of BERT's O activation memory, resp. As shown in Table 2 for the sequence length N = 512, BlOckBERT with 2 / 3 blocks saves 18.7% / resp. The saving more significant for longer sequences. When N 1024, the overall memory reduction of BlOCkBERT with 2/ blocks is 27.3% 36.1%. resp.
23.8% overall memory,
4 EXPERIMENTS
We evaluate the pre-training and fine-tuning performance of BlockBERT. In particular, when = 2, we denote 10:2 to be the configuration which distributes 10 heads to permutation (1, and 2 to permutation (2, 1): when n 3, we denote 8:2:2 to be the configuration which assigns 8, 2, 2 heads
We GPUs of 16 GB memory for profiling. BERT with N 2048 fails due out-of-memory error.

---

Under review as a conference paper at ICLR 2020
N
Model
Training Time (day)
Memory (per GPU, GB)
Heads Config.
Valid, ppl
RoBERTa-
6.62
9.73
3.58
512
BlockBERT n=2
5.83 (-12.0%)
7.91 (-18.7%)
10:2
3,56
BlockBERT n=3
5.80 (-12.5%)
7.32 (-23.8%)
8:2:2
3.71
RoBERTa-Iseq
13.39
3.60
1024
BlockBERT n=2
7.51 (-22.3%)
9.73 (-27.3%)
9:3
3.57
BlockBERT n=3
7.23 (-25.1%)
8,55 (-36.1%)
8:2:2
3.63
Table 2: Pre-training Performance Analysis.
to permutation (1, (2,3, and (3, 1.2,, resp. We compare BlockBERT with the following baselines:
Google BERT The pre-trained base model from Devlin et (2019).
RoBERTa-2seq and RoBERTa-Iseq We compare with two versions of RoBERTa (Liu et al. 2019). RoBERTa-2sse trained with both masked language model (MLM) task and next sentence predic- tion (NSP) task, while RoBERTa-Ised refers to the pre-training model with only MLM task.
SparseBERT We pre-train BERT models with Transformer encoder replaced by a Sparse Trans- former encoder (Child et al.) 2019). We set its sparsity hyper-parameters parameters stride e = 128 and expres- sivity c == 32. The attention masks used for Sparse Transformer encoder are illustrated in Figure
4.1 PRE-TRAINING
the models follow base setting, i.e., L 12, H 768, A and are trained on the same corpus BooksCorpus and English Wikipedia with uncased word piece tokens. We fix the number of tokens per batch x N = 131, 072, i.c., if sequence length N 512 then batch size B 256, 128. The detailed pre-training configuration is listed in Table 6 in Appendix I. Moreover, the pre-training of SparseBERT and BERT follows the RoBERTa- input sequence is up to N tokens until it reaches document boundary. A summary of the pre- training performance comparison between BlockBERT and RoBERTa- Iseq is shown in Table 2 Besides memory saving, we also achieve significant speedup. For example, when N 1024, BlockBERT (n 2) reduces the training time from RoBERTa's 9.7 days 7.5 days.
if sequence length N 1024 then batch size
setting, i.c., we drop the NSP (Next Sentence Prediction) task, and an
4.2 FINE-TUNING TASKS
We evaluate BlockBERT on several question answering tasks, including SQuAD 1.1/2.0 (Rajpurkar
et al. 2018 and five other tasks from the MrQA shared task HotpotQA (Yang et al. 2018), NewsQA Trischler et al. 2017). SearchQA (Dunn et al. 2017 7 TriviaQA (Joshi et al. 200 and
NaturalQA (Kwiatkowski et al. 2019). Since MrQA does not have official test set, we follow
Joshi et al. (2019)) who split the development set evenly to build new development set and test set.
These QA datasets have different paragraph length distribution patterns and are thus ideal for test- ing the effectiveness BloCKBERT. For example, SQuAD, NaturalQA, and HotpotQA consist of mostly short paragraphs (shorter than 512), while paragraphs in SearchQA (average length 1,004) and TriviaQA (average length 934) have around 1,000 tokens. This means that for SearchQA and TriviaQA, a BERT model with sequence length N 512 can only capture half of the context. The detailed paragraph length distributions can be found in Figure 6
For all the pre-trained models, we adopt the same fine-tuning QA setup from Devlin et al. 2009 The tokenized paragraph (P1. ... .P.) and question (q1,-.. 4t) are concatenated to be a sequence [ [SEP]. ). The sequence then fed into the pre-trained model with [CLS] ]1 - q: SEP P1 P.
two extra linear layers for predicting the start and end positions of the answer spans. The detailed
fine-tuning setting is listed in Appendix A.4 Table and Table report the experimental results.
ihups///mraa.itthu.ii

---

Under review as a conference paper at ICLR 2020
SQuAD
1.1
SQuAD 2.0
N
Model
EM
F1
EM
F1
Human Perf.
82.30
91.20
86.80
89.40
Google BERT
81.19
88.45
74.08
77.16
XLNet
78.46
81.33
RoBERTa-2seq
82.91
89.78
75.79
79.17
512
RoBERTa-Isey
84,43
91.48
79.22
82.27
SparseBERT
80.49
88.09
74.15
76.96
BlockBERT n=2, 10:2
84.08
90.77
78.34
81.46
BlockBERT n=3, 8:2:2
82.37
89.64
77.33
80.33
RoBERTa-Iseq
84.58
91.14
79.34
82.26
SparseBERT
81.02
88.37
74.51
77.57
1024
BlockBERT n=2, 9:3
83.65
90.74
78.55
81.45
BlockBERT n=3, 8:2:2
82.74
90.05
76.79
79.84
Table 3: Dev set results on SQuAD 1.1/2.0. The result of XLNet(-Base)) is from Yang al. 2019
BlockBERT (n=2) v.s. RoBERTa- Comparing BlockBERT (n 2) with RoBERTa- Iseq pre-trained model with N = 512, we observe an absolute F1 difference from 0.04 (in NaturalQA) to 1.18 (in NewsQA), with average difference to be 0.55. For N 1024, BlockBERT achieves more comparable or even better performance (in SearchQA, NewsQA, and HotpotQA) to RoBERTa-Iseq. The average difference F1 reduces to 0.27.
BlockBERT V.S. SparseBERT For N = 512, it interesting that BlockBERT with blocks (density 33.33%) %) performs better then SparseBERT (density 44.20%) in both SQuAD and MrQA tasks. Sim- ilar patterns can be observed for N 1024. These results show that off-diagonal masking matrices, e.g., the masking matrix defined by permutation (2, 1), play crucial roles in BlockBERT.
Effeet of Long Sequence Pre-training Our observations are twofold. (1) Long sequence pre- training benefits long sequence fine-tuning. In TriviaQA and SearchQA, of which paragraph lengths are around 1024, pre-training models with N = 1024 achieve significantly better performance. (2) The heterogeneity of pre-training and fine-tuning sequence length may hurt performance. For ex- ample, in SQuAD, we do not see significant performance gain by using pre-trained models with N 1024; in HotpotQA and NewsQA, longer sequence pre-training even hurts performance.
Effect of #Blocks It is not surprising that BlockBERT with 2 blocks (n = 2) performs better than that with blocks (n = 3), because keeps more attention matrix entries. The biggest difference is in SQuAD 2.0 and NewsQA with N 1024, where we observe absolute loss of 1.6 Fl by increasing block number from 2 to
In summary, not only BRO saves training time and memory, but it also has competitive and sometimes better performance, especially for tasks with longer sequences. This demonstrates the effectiveness of our blockwise multi-head attention approach.
SearchQA
TriviaQA EM
NewsQA EM
NaturalQA
HotpotQA EM
N
Model
EM
F1
Fl
Fl
EM
F1
Google BERT
74.94
80.37
70.18
75.35
51.27
66.25
66.13
78.29
60.50
77.08
RoBERTa-2seq
76.12
81.74
71.92
76.79
52.45
66.73
66.98
78.63
61.52
77.81
RoBERTa-Isse Iseq
77.09
82.62
73.65
78.22
56.13
70.64
67.14
79.07
62.77
79.28
512
SparseBERT
73.36
79.01
68.71
73.15
51.18
65.47
65.53
77.46
58.54
74.85
BlockBERT n=2, 10:2
76.68
82.33
72.36
77.53
54.66
69.46
66.94
79.03
62.13
79.15
BlockBERT n=3, 8:2:2
75.54
81.07
72.05
76.74
53.82
68.39
66.14
78.47
60.64
77.46
RoBERTa-Iseg
77.47
83.12
75.29
80.20
55.00
69.64
68.28
80.35
61.89
78.71
SparseBEKT
74.83
80.54
70.56
75.34
51.67
67.16
65.07
77.31
59.65
76.02
1024
BlockBERT n=2, 9:3
77.95
83.57
75.06
79.41
55.44
70.08
67.31
79.39
62.13
78.94
BlockBERT n=3, 8:2:2
76.98
82.76
74.78
79.28
53.48
68.50
65.91
78.20
61.89
78.18
Table 4: MrQA test results (Tasks are sorted decreasingly by average paragraph length).
7

---

Under review as a conference paper at ICLR 2020
4.3 ABLATION STUDY
We fix the assignment of attention heads in above experiments. For example, BlockBERT with sequence length N 512 and blocks is trained with ten heads using permutation (1, and the other two using permutation (2, 1). However, we know that there are other ways to partition twelve attention heads, e.g., seven heads for permutation (1, 2) and the other five for permutation (2, 1). It would be interesting to see how the assignment of heads affects model performance. In this section, we grid search attention head assignments and plot their best validation performance in 1.2M training steps. The results are shown in Figure
Our observations are threefold: (I) Identity permutations, i.c., (1, and (1,2,3), are important. As shown in Figure + all optimal solutions assign considerable attention heads to block diagonal matrices, since those matrices enable each token to attend to its nearby tokens; (2) Non-identity permutations follow the rule of "vital few and trivial many." Although identity permutations important, assigning all attention heads to them (corresponding to 12:0 and 12:0:0 in Figure 44) significantly hurts performance, since the model not learn long-term dependencies with only identity permutation; (3) Pre-training performance and fine-tuning performance are correlated but not always consistent. When n 3, pre-training performance suggests 10:1:1 to be the best head assignment ten heads for permutation (1, 2, 3), one head for (2, and one head for (3, 1)2, but we observe that the configuration of 8:2:2 achieves better performance in fine-tuning tasks.
void pp!
valld toss
voiid pil
mp)
valid
volic
76
155
101
12
18
179
231
577
1011
1902
875
1911
1790
(a) N = 512, 2
(b) N = 1024, 12 =
(c) N = 512,
(d) N = 1024, =3
Figure Ablation over blockwise attention heads assignment.
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. In recent years, there is an increasing interest in training neural networks with low-memory (Sohoni et al. 2019). Mainstream techniques include low-precision training (Micikevicius et al.] 2017), microbatching (Huang et al. 2018). gradient checkpointing (Chen al. (0166). Another line of researches studies this problem from theoreti- cal perspective, including the recently proposed lottery ticket hypothesis (Frankle Carbin 2018)
Since the invention of Transformer (Vaswani al. 2017 Dai et al., 2019) and its successful appli-
cation on language model pre-training (Devlin et al. 2019 Radford et al. 2019 Yang et al. 2019
Liu et al. 2019), there have been several studies attempted to simplify from different perspectives.
The first line of research focuses on attention matrix sparsification, such as Star Transformer (Guo al. 2019). Sparse Transformer (Child al. 2019), Adaptive Sparse Transformer (Correia al. 2019 Sukhbaatar et al.) 2019. Log-Sparse Transformer (Li et al. 2019). etc. However, due to lim- ited support for sparse tensor from current deep learning platforms, most of studies have represent sparse matrix using a dense matrix with a binary mask or rely customized CUDA kernels (Gray al. 2017). The second line of research attempts to prune redundant heads Transformer, such as Voita et al. (2019) and Michel et al. (001)). The third line of research focuses on knowledge dis- tillation, including DistilBERT which distills BERT using a smaller BERT and Tang et al. (2019) which distills BERT with BiLSTM (Hochreiter & Schmidhuber 1997).
8

---

Under review as a conference paper at ICLR 2020
6 CONCLUSION
In this work, we study lightweight BERT model with the goal of achieving both efficiency and effectiveness. We profile and analyze the memory bottlenecks of BERT, and focus optimize dot- product self-attention, which consumes quadratic memory with respect to the sequence length. To reduce both training time and memory consumption, we present BlockBERT, which sparsifies the tention matrices to be sparse block matrices. The proposed model achieves time and memory saving without significant loss of performance. In the future, we would like explore more applications of BlO on NLP tasks involving long sequences such coreference resolution (Joshi et al., 2019b) and document-level machine translation (Miculicich et al. 2018), and also non-NLP tasks such protein sequence modeling (Rives et al. 2019).
REFERENCES
Tianqi Chen, Bing Xu, Chiyuan Zhang. and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:600..0117, 2016.
Rewon Child, Scott Gray, Alec Radford, and lyy Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:194. 2019.
Gonçalo M Correia, Vlad Niculae, and André FT Martins. Adaptively sparse transformers. arXiv 2019.
preprint
Zihang Dai, Zhilin Yang, Yiming Yang, William W Cohen, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint arXiv: 022860, 2019.
Jacob Devlin, Ming- Mei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In NAACL-HLT' 2019, pp. 4171-4186, 2019.
Matthew Dunn, Levent Sagun, Mike Higgins, V Ugur Guney, Volkan Cirik, and Kyunghyun Cho. Searchqa: new q&a dataset augmented with context from search engine. arXiv preprint arXiv: 2017.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint 1XX3:1803300635 2018
Scott Gray, Alec Radford, and Diederik Kingma. Gpu kernels for block-sparse weights. arXiv preprint arXiv:1711.0922. 2017.
Qipeng Guo, Xipeng Qiu, Pengfei Liu, Yunfan Shao, Xiangyang Xue, and Zheng Zhang. Star- transformer. In NAACL-HLT" 2019, pp. 1315-1325, 2019.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint arXiv:1503.022 025331, 2015.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):
1735-1780, 1997.
Yanping Huang, Yonglong Cheng, Dehao Chen, HyoukJoong Lee, Jiquan Ngiam, Quoe V Le, and Zhifeng Chen. Gpipe: Efficient training giant neural networks using pipeline parallelism. arXiv preprint 1rXiv:1811.06665. 2018.
Mandar Joshi, Eunsol Choi, Daniel Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In ACL' 17, pp. 1601-1611, 2017.
Mandar Joshi, Danqi Chen, Yinhan Liu, Daniel Weld, Luke Zettlemoyer, and Omer Levy. Spanbert: Improving pre-training by representing and predicting spans. arXiv preprint /907 10529, 2019a.
Mandar Joshi, Omer Levy, Daniel Weld, and Luke Zettlemoyer. Bert for coreference resolution: Baselines and analysis. arXiv preprint arXiv: 190::00000000 09091, 2019b.

---

Under review as a conference paper at ICLR 2020
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453-466, 2019.
Kenton Lee, Ming- Wei Chang, and Kristina Toutanova. Latent retrieval for weakly supervised open domain question answering. arXiv preprint arXiv: 2019.
Shiyang Li, Xiaoyong Jin, Yao Xuan, Xiyou Zhou, Wenhu Chen, Yu-Xiang Wang, and Xifeng Yan. Enhancing the locality and breaking the memory bottleneck of transformer on time series forecasting. XrXiv preprint arXiv: 1rx.iv/977.0022 002,50 2019.
Yinhan Liu, Myle Out Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907. 11692, 2019.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? preprint arXiv:/905. 10650, 2019.
arXiv
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al. Mixed precision training. arXiv preprint arXh:: 2017.
Lesly Miculicich, Dhananjay Ram, Nikolaos Pappas, and James Henderson. Document-level neural machine translation with hierarchical attention networks. In EMNLP' 18, pp. 2947-2954, 2018.
Rodrigo Nogueira and Kyunghyun Cho.
Passage re-ranking with bert.
arXiv preprint
arXiv:1901. 04085, 2019.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. arXiv preprint arXiv:1904.0103: 2019.
Matthew Peters, Mark Neumann, Mohit Iyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. Deep contextualized word representations. In Proceedings of the 2018 Con- ference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume I (Long Papers), pp. 2227-2237, 2018.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't 't know: Unanswerable questions for squad. arXiv preprint arXiv: 03822, 2018.
Alexander Rives, Siddharth Goyal, Joshua Meier, Demi Guo, Myle Ott, Lawrence Zitnick, Jerry Ma, and Rob Fergus. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. bioRxiv, pp. 622803, 2019.
Livio Baldini Soares, Nicholas FitzGerald, Jeffrey Ling, and Tom Kwiatkowski. Matching the blanks: Distributional similarity for relation learning. arXiv preprint arXiv: (tt:::90600331 2019.
Nimit Sharad Sohoni, Christopher Richard Aberger, Megan Leszezynski, Jian Zhang, and Christopher Ré. Low-memory neural network training: A technical report. arXiv preprint arXiv: 10631, 2019.
Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. arXiv preprint PrXivv1905.07700 2019.
Raphael Tang, Yao Lu, Linqing Liu, Lili Mou, Olga Vechtomova, and Jimmy Lin. Distilling task- specific knowledge from bert into simple neural networks. arXiv preprint arXiv: 1903.126.,, 2019.
10

---

Under review as a conference paper at ICLR 2020
Adam Trischler, Tong Wang, Xingdi Yuan, Justin Harris, Alessandro Sordoni, Philip Bachman, and Kaheer Suleman. Newsqa: machine comprehension dataset. Proceedings of the 2nd Workshop Representation Learning for NLP, pp. 191-200, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads the heavy lifting, the rest can be pruned. arXiv preprint 1oXx::99050094/5 2019.
Zhilin Yang, Peng Qi. Saizheng Zhang, Yoshua Bengio, William W. Cohen, Ruslan Salakhutdinov, and Christopher D. Manning. HotpotQA: A dataset for diverse, explainable multi-ho hop question answering. In EMNLP' 18, 2018.
Zhilin Yang. Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. arXiv preprint arXiv: wwwp:090000.2311 2019.
A APPENDIX
A.1 NOTATIONS AND TRE-TRAINING HYPER-PARAMETER
The notations and pre-training hyper-parameters listed in Table and
and Table
Hyper- parameter
Value
Vocabulary Size
30,522
Dropout
0.1
Description
Base
Large
Attention dropout
0.1
Warmup steps
10K
B
Batch size
256
256
Weight decay
0.01
Self-attention heads
12
Max steps
2.4M
Layers
12
Initial learning rate
0.00025
H
Hidden units
768
1024
Learning rate decay
Linear
4H
Feed-forward hidden units
3072
4096
Adam e
1c-8
N
Sequence length
512
512
Adam 81
0.9
Adam 882
0.999
Gradient Clipping
1.0
Table 5: BERT notations.
Table 6: Pre-training hyper-
parameters.
A.2 PROFILER IMPLEMENTATION
Among the three types of training memory, model memory and optimizer memory is relatively easy to profile (can be computed by enumerate each tenor and summing up tensor numel () tensor. element.ss ze ()) ). To calculate activation memory, [Sohoni et al. (2019) traverse PyTorch's autograd graph and sum up necessary storage space. They find that the summation of model memory, optimizer memory and activation memory matches PyTorch memory profiling tool torch. cuda www.mmxx..r......mmmmmommmmmmmmm www.yyllllll .all Located. Based on their observation, we use
torch. cuda max memory al located model memory optimizer memory
(4)
estimate to activation memory. When profiling BERT, first pre-train it for 1000 steps, and then compute its model and optimizer memory. Finally, we esitmate its activation memory according to Equation
11

---

Under review as a conference paper at ICLR 2020
A.3 SPARSEBERT
The sparse masking matrices we use for Sparse Transformer (Child et al. 2019) are shown in Fig- ure We adopt the implementation from Fairsed
200
200
100
1999
690
(b)
Figure 5: The sparse masking matrices in Sparse Transformer (fixed mode) encoder. White color indicates attention values to masked. (a) N = 512, & = 128, c = 32, density 44.20%; (b) N 1024, 8 = 128, 32, density 34.97%.
A.4 FINE-TUNING SETTINGS
Our fine-tuning is implemented based on code base from HuggingFacq and SpanBERT (Johi al., 2019a). We use max.sequence_uu SaaniceeeenghhhN i.e., we allow fine-tuning task to in- put sequences as long as the pre-training model. If the input sequence is too long to fit the max_sequence._eng-te
constraints. we use sliding window of stride 128 split it. We grid search learning rate from {5c-6, 1c-5, 2c-5, 3e-5, 5e-5} } and batch size from {16, 32}. The fine-tuning performed for 4 epoches.
(183)
AweiEA PanumalAA (361
HurgezA (218
(534)
SQuAD (13)
Histogram
0.0%
000
Paragraph Length
Figure 6: Paragraph length (after tokenization) distribution. The distribution of SQuAD 2.0 is very similar to SQuAD 1.1, so we only plot SQuAD 1.1 here.
A.5 TEST EFFICIENCY
We benchmark efficiency of RoBERTa and proposed BlockBERT. The benchmark code follows huggingface 10 All experiments are run 30 times on a 32GB V100 GPU with half pre- cision (FP16). We report the average running time at Table As we can see. BlockBERT does achieve speedup and memory reduction during test time. Take 8x 1024, i.c., batch size B = 8, sequence length
1024, as example, we can see that BlockBERT with blocks saves 27. 27.8% of test time, and BlockBERT with 3 blocks saves more (30.4%). As for memory, we can observe that RoBERTa can not handle an input 16x 1024, while possible for BlockBERT to work on it.
hhtprr/giiiiiaoooooooeooeooooeei multihh attention.py.
12

---

Under review as a conference paper at ICLR 2020
B x N
8x 1024
16x 1024
24 1024
32x 1024
RoBERTa
0.1371 0.0990 0.0954
OOM
OOM
OOM OOM OOM
BIOC n=2 BlockBERT n=3
0.1869 0.1790
OOM 0.2634
Table 7: Test time statistics (sec) for different input size (Batch size x Sequence length). OOM indicates out-of out-of-memoon off-meooy.
13


Reviews:
[{"review_id": "SkgHIb4G9B", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose BlockBERT, a model that makes the attention matrix of Transformer models sparse by introducing block structure. This has the dual benefit of reducing memory and reducing training time. The authors show on various question answering tasks that their model is competitive with RoBERTa. \n\n1. Can the authors add additional details about the training of their model in Section 4.1? It is not clear for example the vocabulary size - is that the RoBERTa vocabulary size of around 50k or the BERT vocabulary size that is smaller? I believe this will affect the memory and training speed. \n\n2. For the datasets such as TriviaQA and SearchQA, how is RoBERTa finetuned on these tasks? By doing the window approach? \n\n3. The authors compare to RoBERTa and Sparse BERT as baselines for the section on performance. However, can the authors also include metrics on training time and memory in Table 2 for Sparse BERT as well as other sparse attention transformer architectures proposed (for example the Correia paper or the Sukhbaatar paper)? It is not clear the savings from this architecture compared to sparse Transformers in general. \n\n4. The analysis in Section 4.3 is quite unclear due to how compressed the description is and how tiny the graphs are. \n\n5. The authors mention that attention heads can be sparsified due to the memory usage and quadratic time complexity. Other work has also shown that the attention heads are quite redundant and can be pruned away, so attention head dropout is effective. For example https://arxiv.org/abs/1905.09418 \n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "SJlkRGil9B", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper introduces a optimisation for BERT models based on using block matrices for the attention layers. This allows to reduce the memory footprint and the processing  time during training while reaching state-of-the-art results on 5 datasets. An interesting study on memory consumption in BERT is conducted. No results are given at test time : is there also a memory and processing time reduction ?\n\nEven if the proposition is interesting, the impact of the paper is limited to the (flourishing) scope optimising Bert models (\"Bertology\"). The authors do not mention if their code is available.\n\n\nTable 3 : Humam -> Human\t", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "HJgqUsBatS", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes blockwise masked attention mechanisms to sparsify Transformer architectures, the main motivation being reducing the memory usage with long sequence inputs. The resulting model is called BlockBERT. The paper falls in a trend of recent papers compressing/sparsifying/distilling Transformer architectures, a very relevant area of research given the daunting resources needed to train these models.

While the proposed contribution is very simple and interesting, it also looks a rather small increment over prior work, namely Sparse Transformer and Adaptive Span Transformer, among others. Experiments are rather limited and the memory/time reduction is not overwhelming (18.7-36.1% less memory, 12.0-25.1% less time), while final accuracy is sometimes sacrificed by a few points. No comparison to other adaptively sparse attention transformer architectures (Correia et al. EMNLP 19 or Sukhbaatar et al. ACL 19) which should as well provide memory reductions due to the sparsity of the gradients, which require less activations to be cached. I suggest addressing this concerns in an eventual resubmission of the paper.
Decision: Reject