Block Skim Transformer for Efficient Question Answering
Authors: Yue Guan
Jingwen Leng
Yuhao Zhu
Minyi Guo
Venue:  (2021)
Keywords: Efficient Transformer
Question Answering

Abstract:   Transformer-based encoder models have achieved promising results on natural language processing (NLP) tasks including question answering (QA).  Different from sequence classification or language modeling tasks, hidden states at all positions are used for the final classification in QA.  However, we do not always need all the context to answer the raised question.  Following this idea, we proposed Block Skim Transformer (BST) to improve and accelerate the processing of transformer QA models.  The key idea of BST is to identify the context that must be further processed and the blocks that could be safely discarded early on during inference. Critically, we learn such information from self-attention weights.  As a result, the model hidden states are pruned at the sequence dimension, achieving significant inference speedup.  We also show that such extra training optimization objection also improves model accuracy. As a plugin to the transformer-based QA models, BST is compatible with other model compression methods without changing existing network architectures.  BST improves QA models' accuracies on different datasets and achieves $1.6\times$ speedup on $BERT_{large}$ model.


Under review as a conference paper at ICLR 2021
BLOCK SKIM TRANSFORMER FOR EFFICIENT QUESTION ANSWERING
Anonymous authors Paper under double-blind review
ABSTRACT
Transformer-based encoder models have achieved promising results on natural language processing (NLP) tasks including question answering (QA). Different from sequence classification or language modeling tasks, hidden states at all positions are used for the final classification in QA. However, we do not always need all the context to answer the raised question. Following this idea, we proposed Block Skim Transformer (BST) to improve and accelerate the processing of transformer QA models. The key idea of BST is to identify the context that must further processed and the blocks that could be safely discarded early on during inference. Critically, learn such information from self-attention weights. As result, the model hidden states are pruned at the sequence dimension, achieving significant inference speedup. We also show that such extra training optimization objection also improves model accuracy. As a plugin to the transformer based QA models, BST is compatible other model compression methods without changing existing network architectures. BST improves QA models' accuracies on different datasets and achieves 1.6x speedup BERT model.
INTRODUCTION
With the rapid development of neural networks in NLP tasks, the Transformer (Vaswani al., 2017) that uses multi-head attention (MHA) mechanism is recent huge leap (Goldberg, 2016). It has become standard building block of recent NL.P models. The Transformer-based BERT (Devlin et al., 2018) model further advances the model accuracy by introducing self-supervised pre-training and has reached the state-of-the-an accuracy on many NLP tasks.
One of the most challenging tasks in NLP is question answering (QA) (Huang et al., 2020). Our key insight is that when human beings are answering a question with a passage as context, they do not spend the same level of comprehension for each of the sentences equally across the paragraph. Most of the contents are quickly skimmed over with little attention on it. However, in the Transformer architecture, all tokens go through the same amount of computation, which suggests that we can take advantage of that by discarding many of the tokens in the early layers of the Transformer. This redundant nature of the transformer induces high execution overhead on the input sequence dimension.
To mitigate the inefficiencies in QA tasks, we propose assign more attention some blocks that are more likely to contain actual answer while terminating other blocks early during inference. By doing so, we reduce the overhead of processing irrelevant texts and accelerate the model inference. Meanwhile, by feeding the attention mechanism with the knowledge of the answer position directly during training, the attention mechanism and QA model's s accuracy are improved.
In this paper, we provide the first empirical study on attention featuremap to show that an attention map could carry enough information to locate the answer scope. We then propose Block Skim Transformer (BST), a plug-and-play module to the transformer-bat based models, to accelerate transformer- models on QA tasks. By handling the attention weight matrices as feature maps, the CNNN-bsed Block Skim module extracts information from the attention mechanism to make skim decision. With the predicted block mask, BST skips irrelevant context blocks, which do not enter subsequent layers' computation. Besides, devise new training paradigm that jointly trains the Block Skim

---

Under review as a conference paper at ICLR 2021
objective with the native QA objective, where extra optimization signals regarding the question position are given to the attention mechanism directly.
In our evaluation,
show BST improves the QA accuracy and Fl score all the datasets and models we evaluated. Specifically, BERT is accelerated for 1.6x without any accuracy loss and nearly 1.8x with less than 0.5% F1 score degradation.
This paper contributes to the following 3 aspects.
We for the first time show that an attention map is effective for locating the answer position in the input sequence.
We propose Block Skim Transformer (BST), which leverages the attention mechanism to improve and accelerate transformer models on QA tasks. The key is to extract information from the attention mechanism during processing and intelligently predict what blocks to skim.
We evaluate BST several Transformer-basee model architectures and QA datasets and demon- strate BST efficiency and generality.
2 RELATED WORK
Recurrent Models with Skimming. The idea to skip skim irrelevant section tokens of input sequence has been studied NLP models, especially recurrent neural networks (RNN) (Rumelhart et al., 1986) and long short-term memory network (LSTM) (Hochreiter & Schmidhuber, 1997). LSTM-Jump (Yu et al., 2017) uses the policy-gradinn graddent reinforcement learning method to train a LSTM model that decides how many time steps jump at each state. They also hyper-parameten to control the tokens before jump, maximum tokens to jump, and maximum number of jumping. Skim-RNN (Seo et al., 2018) dynamically decides the dimensionality and RNN model size to be used at next time step. In specific, they adopt two "big" and "small" RNN models and select the "small" one for skimming. Structural-Jump-LSTT jump small step to next token or structurally to next punctuation. Rkip-RNN (Campos et al., 2017) learns to skip state updates thus results in reduced computation graph size. The difference of BST to these works are two-fold. Firstly, the previous works make skimming decisions based on the hidden states or embeddings during processing. However, we are the first to analyze and utilize the attention relationship for skimming. Secondly, our work is based on Transformer model (Vaswani et al., 2017), which has outperformed the recurrent type models on most NLP tasks.
(Hansen et al., 2018) use two agents decide whether
Transformer with Input Reduction. On contrast to aforementioned recurrent models, in the pro- cessing of Transformer- based model, all input sequence tokens are calculated in parallel. As such, skimming be regarded reduction sequence dimension. PoWER-BERS (Goyal al., 2020) extracts input sequence token-wir wise during processing based on attention scores to each token. During the fine-tuning process for downstream tasks, Goyal al. proposes soft-extract layer train the model jointly. fuunel--ransformer (Dai et al., 2020) proposes a novel pyramid architecture with input sequence length dimension reduced gradually regardless of semantic clues. For tasks requiring full sequence length output, like masked language modeling and extractive question answering, Funnel-Transformer up-sample the input dimension to recover. Universal Transformer (Dehghani et al., 2018) proposes dynamic halting mechanism that determines the refinement steps for each token. Different from these works, BST utilizes attention information between question and token pairs and skims the input sequence at the block granularity accordingly.
Efficient Transformer. There are also many attempts for designing efficient Transformers (Zhou et al., 2020; Wu et al., 2019; Tay et al., 2020). Well studied model compression methods for Transformer models include pruning (Guo et al., 2020), quantization (Wang & Zhang, 2020), distil- lation (Sanh et al., 2019), weight sharing. Plenty of works and efforts focus on dedicated efficient attention mechanism considering its quadratic complexity of sequence length (Kitaev et al., 2019; Beltagy al., 2020; Zaheer et al., 2020). BST is orthogonal to these techniques the input dimen- sion and therefore is compatible with them. We demonstrate this feasibility with the weight sharing model Albert (Lan al., 2019) in Sec. 5.
2

---

Under review as a conference paper at ICLR 2021
1.0
Logistic Regrossion
0.8
0.6
Accuracy
12
16
20
24
Layer
Figure 1: Accuracy and score performance logistic regression model predicting whether a block contains answer. The logistic regression model built with attention values between question and larget block as input feature.
PROBLEM FORMULATION: Is ATTENTION EFFECTIVE FOR SKIM
Transformer. Transformer model with multi-head self-attention mechanism calculates hidden states for each position weighted sum of input hidden states. The weight vector is calculated by parameterized linear projection query Q and key as eq. 1. Given sequence of input embeddings, the output contextual embedding composed by the input sequence with different attention each position.
QKT Vig
Attention..x Softmax
where K are query and key matrix of input embeddings, dg is the length of a query or key vector. Multiple parallel groups of such attention weights, also referred attention heads, make possible to attend to information at different positions.
QA is one of the ultimate downstream tasks in the NLP. Given text document and question about the context, the answer is a contiguous span of the text. To predict the start and end position of the input context given question, the embedding of each certain token processed for all transformer layers in the encoder model. In many end-to-end open domain QA systems, information retrieval the advance procedure at coarse -rriined passage or paragraph level. Under the characteristic of extractive QA problem that answer spans are contiguous, our question is that whether we can utilize such idea at fine-grained block granularity during the processing of transformer. the attention weights effective for distinguish the answer blocks?
To answer the above question, we build a simple logistic regression model with attention matrix from each layer predict whether input sentence block contains the answer. The attention matrices profiled from BERT SQuAD QA model and reduced to block level following Eq. (Clark et al., 2019). The attention from block [a.b] attending to block [c,d] is aggregated to value. And the attention between a block and the question sentence, special tokens [CLS] and [SEP] " are used to denote the attending relation of the block. Such 6-dimensional vector from all attention heads in the layer are concatenated as the final classification feature. The result is shown in Fig. with attention matrices from different layers. Simple logistic regression with hand crafted feature from attention weight achieves quite promising classification accuracy. This suggests that the attending relationship between question and targets indeed capable for figuring out answer position.
(luckAttention((ubbieee a i
Σ
Allention(i.])
3

---

Under review as a conference paper at ICLR 2021
(c]Block Skim Module
(b) Block Skim Transformer
(a) Training Objective
LOSSeni
bbbpee Transformer Layeri
Block Skin Midule
Lineet
Biock Mask
LotgA
Losiig
QA Prediction
Block Module
Convi
Attention Featuramop
nottnniion Herdy
Block
A RST,
Trankemm Layan
Convara
Block Mask
AvgPooling"?
Tersformen
BST,
Boitch Notrmmlatinn
Convaty
Boock Moskk
AvgPooling22
000000
0
Input Sequence
shape-[n.AH]
Figure 2: (a) The overview of Block Skimming Transformer. (b) A BST layer is composed of a Transformer layer and Block Skimming Module. (C) The CNN-based BSM module details.
4 BLOCK SKIMMING TRANSFORMER (BST)
4.1 ARCHITECTURE OVERVIEW OF BST
We propose the Block Skimming Transformer (BST) model to accelerate the question answering task without degrading the answer accuracy. Unlike the conventional Transform-band baseed model that uses all input tokens throughout the entire layers, our BST model accurately identifies the irrelevant contexts for the question in the early layers, and remove those irrelevant contexts in the following layers. As such, our model reduces the computation requirement and enables fast question answering.
In Sec. 3, we have shown that it is feasible to identify those tokens that are irrelevant to the question through a hand-crafted feature using the attentions relationship among tokens. However, using this approach could significantly hurt the question answering task accuracy as we show later. As such, we propose end-to-end learnable feature extractor that captures the attention behavior better.
Fig. 2 shows the overall architecture of our BST model, where layer is composed of a Transformer layer and learnable Block Skim Module (BSM). The BSM adopts the convolutional neural network for feature extraction. The input is attention matrices of attention heads, which are treated as feature maps of multiple input channels. The output is block-level mask that corresponds the relevance of a block input tokens to the question.
In each BSM module, reduce the size of feature maps. Two x convolution and one x convolution are connected with pooling operations intersected. For all the convolution operations, ReLU funcition (Hahnloser & Seung, 2001) is used as activation function. To locate the answer context blocks, we use a linear classification layer to calculate the score for each block. Also, Batch Normalization layers (Ioffe & Szegedy, 2015) are inserted to improve the model accuracy.
convolution to collect local attending information and pooling to
Formally, we denote the input sequence of transformer layer as (Xo, Then the atten- tion matrices of this layer are denoted as Attention(X), Given the attention output of transformer layer, the kn block prediction result represented as B = BST (Artention(x)), where BST the proposed architecture. The main functions of BST is expressed as Eq. 3.
BST (Attention) Lineer((Conv (Convaxa (Pool (Polansss(((l(AAAAeetm)))))))))
(3)
4.2 JOINT TRAINING OF QA AND BLOCK-SKIN CLASSIFIERS
There types of classifiers in our BST model. where the first is the original QA classifier the last layer and the second the block-level relevance classifier at each layer. We jointly train these classifiers that the training objective is minimize the sum of all classifiers' losses.
4

---

Under review as a conference paper at ICLR 2021
The loss function of each block-level classifier is calculated as the cross entropy loss against the ground truth label whether a block contains answer tokens or not. Equation 4 gives the formal definition. The total loss of the block-level classifier LesT is the sum of all blocks that only contain passage tokens. The reason is that we only want to throw away blocks with irrelevant passage tokens instead of questions. Blocks that have question tokens or padding tokens are not used in the training process. To be more detailed, such blocks are pre- see-processed and dropped during the training process.
LBST
Σ
CELoss(b,.))
be(passage blocks)
(4)
1 block i has answer tokens
y1
0 block i has no answer tokens
To calculate the final total loss Lrota,, we introduce two hyper-parameters in Equation 5. We first use the syper-parameter a so that different models and settings could adjust the ratio between the QA loss and block-level relevance classifier loss. We then use the other hyper-parameter & to balance the loss from positive and negative relevance blocks because there are typically many more blocks that contain no answer tokens (negative bocks) than the blocks that do contain answer tokens (positive bocks). We explain how tune those ryper-parameters for different models and settings later.
Liotal = Les a Σ
(B BST
-oly- BST
(5)
4th layer
Although we add the block-level relevance classification loss in the joint training, we do not actually throw away any blocks because it can skip answer blocks and the QA task training becomes unstable. In this sense, the block-level relevance classification loss can be viewed as a regularization method for the QA training as we force attention heads to better distinguish the answer blocks and non-answer blocks. As we show in the experiment, this regularization effect leads accuracy improvement for the QA task.
4.3 USING BST FOR QA
We now describe how to the BST model to accelerate the QA task. In the above joint training process, we add the BSM module in every layer. However, we only augment a specific layer with the BSM module during the inference to save computation and avoid heavy changes the underlying Transformer model. As such, the layer index for augmenting is a hyper-parameted in our model.
Once the BSM-augmented layer is chosen, split the input sequence by the block granularity, which is another byper-parameter in our model. The model skips a set of blocks according to the BSM module results for the following layers. It should be noted that the BST training process does not throw away any blocks because if a relevant block with answer tokens is rejected, the training of the original QA task confused and becomes unstable.
To maintain compatibility with the original Transformer model, we forward the skipped blocks directly to the last layer for the QA classifier. With those design features, BST works as an add-on component to the original Transformer model and is compatible with many Transformer variant models as well as model compression methods. specific, we will demonstrate that BST works well with Transformer-basee Roberta (Liu et al., 2019), which has different pre-training objective and sequence encoding, and Albert (Lan et al., 2019), which shares weights among layers for a reduced model size.
We provide an analytical model to demonstrate the speedup potential of BST Suppose that we insert the BSM module at the layer out of the total L layers, and portion of k blocks remain for the following layers. The performance speedup is formulated by Equation if we ignore the computation overhead in the BSM module. In fact, the computation of single BSM module is much smaller than Transformer layers. For example, if k 1/3 blocks remain after I/L = 1/3 layers, we achieve a 1.8x ideal speedup. Similarly, if 1/4 blocks remain after I/L 1/4 layers, the ideal speedup is 2.29
L.N Thayer
speedup
(6)
N I.N Tuayer + (L- 1). N..KK Tiayer 1- (1-1111)(11
5

---

Under review as a conference paper at ICLR 2021
QA
BST classifier
baseline
BST
layer 4
middle layer
EM
F1
EM
Fl
F1
Fl
Transformer-basee models with BST on SQuAD
BERT base
81.14
88.52
81.38
88.69
80.99
89.37
BERT large
86.53
92.82
87.12
93.20
82.48
93.32
Robertabase
82.02
89.65
82.40
89.91
81.88
90.35
Robertatarge
86.03
92.81
86.17
92.95
78.93
86.56
Alberthase
82.13
89.78
82.18
89.78
84.30
90.25
A Alberttarge
84.26
91.51
84.71
91.73
78.93
93.02
Avg.
83.68
90.85
83.99
91.04
81.38
90.48
BERT base
BST on QA datasets
SQuAD
81.14
88.52
81.38
88.69
80.99
89.37
NewsQA
51.45
66.57
52.28
67.43
53.74
67.02
TriviaQA
68.99
73.78
69.04
73.91
93.96
96.20
HotpotQA
58.76
75.56
59.55
75.93
69.67
75.36
Natural Questions
67.00
79.00
67.47
79.15
89.32
90.54
Avg.
65.47
76.69
65.94
77.02
77.54
83.70
Table I: BST training results. (Upper) BST evaluated with different Transformer variant models and model sizes. (Lower) BST evaluated with BERT base on 5 QA datasets.
EXPERIMENT
5.1 EXPERIMENTAL SETUP
Dataset. We evaluate our method on 5 extractive QA datasets, including SQuAD 1.1 (Ra- jpurkar et al., 2016), Natural Questions (Kwiatkowski al., 2019), TriviaQA (Joshi et al., 2017), NewsQA (Trischler et al., 2016) and HotpotQA (Yang et al., 2018). The diversity of these datasets such as various passage lengths and different document sources lets us evaluate the general appli- cability of the proposed BST method. We follow the setting of BERT model use the structure of Transformer encoder and a linear classification layer for all the datasets.
Model. As mentioned earlier, BST works plugin module to the oracle Transformer model, and therefore applicable to Transformer- based models. To illustrate this point, we apply our method three different models including BERT, Roberta (Liu et al., 2019) with a different pre-training objective, and Albert (Lan et al., 2019) with parameter sharing layers. For all three models, we evaluate the base setting with 12 heads and 12 layers and the large setting with 24 layers and 16 heads as described in prior work (Devlin et al., 2018).
Training Setting, We implement the proposed method based on open-sourced library from Wolf et al. (2019). For each baseline model, we use the released rre-trained checkpoints We follow the training setting used by Devlin al. (2018) and Liu al. (2019) to perform the fine-tuning on the above extractive QA datasets. We initialize the learning rate to 3e 5 for BERT and Roberta and Se for Albert with linear learning rate scheduler. For SQuAD dataset, we apply batch size 16 and maximum sequence length 384. And for the other datasets, we apply batch size 32 and maximum sequence length 512. We perform all the experiments reported with random seed 42. We train a baseline model and BST model with the same setting for epochs and report accuracies from MRQA task benchmark for comparison. We use four V100 GPUs with 32 GB memory for training and report performance speedup multiple different hardware platforms.
For the following experiments, we use the block size 32 unless explicitly mentioned. We set the hyper-parameter ß to 4 for all experiments and to 1 except Albert. We the value of 0.05 for Albert. In the Albert model, the parameters of transformer layers are shared but BST modules
We use pre-trained language model checkpoints released from https: htahee////gggggggg comode11

---

Under review as a conference paper at ICLR 2021
(a) BERT base
(b) BERT large
2.0
1.35
1.8
1.30
Speedup
Spoodup
1.6
1.20
1.15
1.4
1.10
baseline
12
layer
layer
layer
1.00
layer 6
1.0
layer 12
87.25
87.50
87.75
8B
88.50
88.75
91.0
91.5
92.0
92.5
93.0
F1
Figure 3: Speedup of BST on BERT hase and BERT large augmented at different layers.
in our method do not share parameters. As such, we decrease the loss from BST to prevent model over-fitting and its impact on the QA task parameters.
Performance Evaluation. We measure the performance speedup of our method on the 2.20GHz Intel Xeo Silver 4210 CPU with the batch size of one. On the GPU, the batch size of one could not fully utilize the GPU computation resource so the inference time is bottlenecked at the memory. Our evaluation scenario closely resembles prior work (Wu et al., 2019) that targets the mobile application domain. We evaluate BST with difffernt layers and prediction thresholds of BST classifier to explore the trade-off between performance speedup and model accuracy. For example, a lower prediction threshold could lead to more skipped blocks, which means better performance speedup. On the other hand, also increases the chance of skipping answer blocks, which hurts the QA task accuracy.
5.2 BST AS A REGULARIZATION МETHON
We first evaluate BST model as a regularization method to improve the QA task accuracy. Specifically, we compare the accuracy of three baseline models and their BST variants. In their BST versions, the BSM modules only participate the training process, and are removed in the inference task.
The upper half Tbl. shows the accuracy comparison on SQuAD dataset. By only changing the training process, BST improves the extractive QA accuracy for all baseline models. On average, BST exceeds the baseline by 0.32% in exact match and by 0.19% in Fl score. The accuracy improvement of BST is generally greater on large models. We attribute this to a stronger regularization effect for larger models. The results show the wide applicability of our method to different models.
Tbl. also demonstrates the BST classifier F1 score trained jointly (but not used in this setting) with the baseline models. For simplicity, we only show the results of layer 4 and the middle layer (layer for base and 12 for large model). On average, the block-level relevance classifier has notably high F1 score even at early layers (averaged 81.38%) and even higher scores at the middle layer.
The lower half of Tbl. shows BERT large results multiple QA datasets. BST outperforms the baseline training objective on all datasets evaluated and exceeds with 0.52% exact match and 0.33% Fl score on average. The results show the wide applicability of our method different datasets with varying difficulty and complexity. Meanwhile, we also observe modest correlation between the block-level relevance classifier and the QA task. In other words, the BST classifier tends be higher on datasets with higher QA accuracy except for TriviaQA dataset.
5.3 QA TASK SPEEDUP WITH BST
We now demonstrate the BST's ability to accelerate the QA task. Fig. 3 demonstrates the performance speedup against Fl score evaluated with BERT and BERTharg model on SQuAD dataset. By tuning the prediction threshold of the BST classifier, we can trade-off between acceleration and accuracy loss. Here evaluate the BST classifier with 0.9.0.09. 0.99 prediction threshold with
7

---

Under review as a conference paper at ICLR 2021
Update Transformer
Skim Training
Augment Layer
Block Size
QA
BST Classifier
id
BSM
EM
FI
layer 4
layer
1
Bascline
81.14
88.52
2
Vanilla BST
All
32
81,38
88.69
80.99
89.37
3
Freeze Transformer
All
32
81.14
88.
79.60
78.41
Hand-crafted Feature
81.21
88.59
69.41
69.51
5
Augment One Layer
81.20
88.63
77.26
6
Skim Traning
79.27
86.83
84.76
7
Block Size
81.22
88.60
75.83
85.03
8
Block Size 8
All
81.25
88.63
76.
87.04
9
Block Size 16
AlI
16
81.35
88.75
78.41
82.51
10
Block Size 64
All
64
81.39
88.65
87.77
90.76
11
Block Size 128
128
80.90
88.33
91,36
92.79
Table 2: Ablation studies of the BST design components of BERT base on SQuAD dataset.
classifier at layer and the middle layer respectively. On BERT base, BST achieves 1.38x speedup with the same accuracy to the baseline. With more aggressive skipping strategy, 1.4x speedup is obtained with minor accuracy loss (less than 1.5%). On BERT large BST achieves 1.6x speedup with minor accuracy improvement and nearly 1.8x speedup with less than 0.5% F1 score degradation.
Generally, the specific layer for inserting the BSM module can be determined by hyper-parameter search according Eq. 6. As shown in Fig. 3, skipping the irrelevant blocks layer tends to be better than that at the middle layer, which is layer 6 for BERT base and layer 12 for BERT This is because more computation is reduced when skipping earlier layers and the BST classifier already has a quite good prediction accuracy at early layers.
5.4 ABLATION STUDY
compare our BST method with a series ablation design components to study their individual effect. The experiments are performed based the same setting as Sec. 5.1. We perform the experiments described in Tbl. 2, which has also the detailed results, and summarize the key finds as follows.
(3) Instead of joint training as described in Sec. 4.2, we perform two-step training. We first perform the fine-tuning for the QA task. We then perform the BSM module training with the baseline QA model frozen. In other words, we only use the BST objective and only update the weights in the BSM modules. Therefore, the QA accuracy remains the same as the baseline model, which is lower than the joint training (id 3). Meanwhile, the BST classifier also has a lower accuracy than the joint training especially at layer 6.
(4) Instead of BSM module, we use the hand-crafted feature in Sec. 3. The resulted block-level relevance classification accuracy considerably lower than our learned BST model.
(5) Instead of adding BST module to all layers, we only deploy it into one layer. The experiment result shows that it is beneficial for the model to have BST loss added to every layer.
(6) We skim blocks during the joint QA-BST training process. Because the mis-skimmed blocks may confuse the QA optimization, this training strategy results considerable accuracy loss.
(7-11) We evaluate the accuracy with different block sizes. Specifically, when the block size is it is equivalent to skim at the token granularity. Our experimental result shows that the accuracy of BST classifier is better when the block size is larger. On the other hand, a larger block size also leads to less number of blocks and therefore the performance speedup becomes limited on the studied datasets. To this end, we choose the block size of 32 as a design sweet spot.
CONCLUSION
In this work, we provide plug-and-pp module BST to Transformer and its variants for efficient QA processing. Our empirical study shows that the attention mechanism in the form of a weighted feature map can provide instructive information for locating the answer span. In fact, we find that attention
8

---

Under review as a conference paper at ICLR 2021
map can distinguish between answers and other tokens. Leveraging this insight, we propose to learn the attention in a supervised manner. In effect, BST terminates irrelevant blocks at early layers, significantly reducing the computations. Besides, the proposed BST training objective provides attention mechanism with extra learning signal and improves QA accuracy on all datasets and models evaluated. With the use of BST module, such distinction is strengthened in supervised fashion. This idea may be also applicable to other tasks and architectures.
REFERENCES
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint XXiv:2004.051500, 2020.
Victor Campos, Brendan Jou, Xavier Giró-i-Nicto, Jordi Torres, and Shih-Fu Chang, Skip RNN; learning skip state updates in recurrent neural networks. CoRR, 0tss/770..0881. 2017. URL http: /arxiviivvv/////////// orgg wtwg/aa//////oooo..e.../oe.eaeee 11008 06834.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D Manning. What does bert look at? an analysis of bert's attention. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyyin and Interpreting Neural Networks for NLP, pp. 276-286, 2019.
Zihang Dai, Guokun Lai, Yiming Yang, and Quoc Le. traneortranssormmrr: Filtering out sequential redundancy for efficient language processing. arXiv preprint arXiv:2006.0322. 2020.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations, 2018.
Jacob Devlin, Ming-Was Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep
bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805. 2018.
Yoav Goldberg. A primer neural network models for natural language processing. Journal of Artificial Intelligence Research, 57:345 4200 2016.
Saurabh Goyal, Anamitra Roy Choudhary, Venkatesan Chakaravarthy, Saurabh ManishRaje, Yogish Sabharwal, and Ashish Verma. Power-bert: Accelerating bert inference for classification tasks. arXiv preprint arXhv:2001.08950 2020.
Cong Guo, Bo Yang Hsueh, Jingwen Leng, Yuxian Qiu, Yue Guan, Zehuan Wang, Xiaoying Jia, Xipeng Li, Minyi Guo, and Yuhao Zhu. Accelerating sparse dnn models without hardware-suppon via tile-wise sparsity. arXiv preprint arXiv:2008. 13006,, 2020.
Richard HR Hahnloser and H Sebastian Seung. Permitted and forbidden sets in symmetric threshold- linear networks. In Advances in neural information processing systems, pp. 217-223, 2001.
Christian Hansen, Casper Hansen, Stephen Alstrup, Jakob Grue Simonsen, and Christina Lioma. Neural speed reading with druucural-jump-bee In International Conference Learning Repre- sentations, 2018.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1785-1780, 1997.
Zhen Huang, Shiyi Xu, Minghao Hu, Xinyi Wang, Jinyan Qiu, Yongquan Fu, Yuncai Zhao, Yuxing Peng, and Changjian Wang. Recent trends in deep learning based open-domain textual question answering systems. IEEE Access, 8:94341-94355 2020.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, pp. 448-456, 2015.
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 1601- 1611, 2017.

---

Under review as a conference paper at ICLR 2021
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2019.
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.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel. Piyush Sharma, and Radu Soricut. Albert: A lite bert for self- self-supe learning of language representations. In International Conference on Learning Representations, 2019.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. 2019.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100, 000+ questions for machine comprehension of text. In EMNLP, 2016.
David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by back-propagating errors. nature, 323(6888):533-556, 1986.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108. 2019.
Minjoon Seo, Sewon Min, Ali Farhadi, and Hannaneh Hajishirzi. Neural speed reading via skim-rnn. In International Conference on Learning Representations, 2018.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. arXiv e-prints, pp. arXiv-2009, 2020.
Adam Trischler, Tong Wang, Xingdi Yuan, Justin Harris, Alessandro Sordoni, Philip Bachman, and Kaheer Suleman. Newsqa: A machine comprehension dataset. 2016.
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-50098 2017.
Chunpei Wang and Xiaowang Zhang. Q-bert: A bert-based framework for computing sparqi similarity in natural language. In Companion Proceedings of the Web Conference 2020, pp. 65-66, 2020.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac. Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Huggingface's transformers: State-of-the-au natural language processing. ArXiv, abs/1990.03371, 2019.
Zhanghao Wu, Zhijian Liu, Ji Lin, Yujun Lin, and Song Han. Lite transformer with long-short range attention. In International Conference on Learning Representations, 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-hop question answering. In EMNLP, 2018.
Adams Wei Yu, Hongrae Lee, and Quoc Le. Learning to skim text. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp- 1880-1890, 2017.
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, al. Big bird: Transformers for longer sequences. arXiv preprint arXiv:2007. 14062, 2020.
Wangchunshu Zhou, Canwen Xu, Tao Ge, Julian McAuley, Ke Xu, and Furu Wei. Bert loses patience: Fast and robust inference with early exit. arXiv, pp. arXiv-2006. 2020.
10


Reviews:
[{"review_id": "xpYtNdYUsF7", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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 block skimming approach for QA. \n\nThe main idea is to use a light-weight module that uses the attention weights to predict which blocks of input are not useful for answering a question and use this information to avoid processing these blocks in further layers. Empirical evaluation shows that this approach can reduce computation and provide as much as 1.6X speedup with no drop in accuracy and even further speedups for more small drops in accuracy. Furthermore, using the block skimming component i.e. recognizing which blocks are useful works as a regularizer and can even increase QA performance by modest amounts. \n\nThe main strengths of the work are: \n\n1. The main idea is interesting and well-motivated. We want good solutions for getting large models to run efficiently.\n2. The evaluation on multiple datasets and across three types of transformer based models show that skimming can provide modest but consistent gains and can provide speedups.\n\nThe main weaknesses are: \n\n1. Efficiency is a key point of this method. Yet, there are no external efficiency related baselines compared (e.g. DistillBERT). I buy the argument that this method can be seen as orthogonal to the other methods but it will be useful at the very least to show how the speed-ups fare when compared to other methods. The only comparison we have is with Albert but we don’t have any speed-up comparisons with it either.  \n2. The gains with using skimming as a regularization method is limited. \n\n\nOther Questions and Suggestions:\n\n\n1. Given the limited improvements in effectiveness, it will be worth reporting averages of a few runs (three-five) for at least  some of the rows in Table 1. \n\n2. It will be useful to know the speed-ups for some of the block sizes. Is it possible to add these numbers in Table 2? \n\n3. DeFormer [1], a paper that appeared in ACL 2020, is a related work that simply removes attention computation across question and passage blocks in lower layers. This idea here is more general than the DeFormer work but is close enough to warrant a discussion in the paper at the very least.\n\n4. Can you provide some intuitions for the specific design choices for the Block Skimming Module? \n\n5. It will be interesting to show if you are able to run BERT-large in the skim mode faster than  BERT-base in the original model. This will highlight the fact that skimming method can be used to improve\n\n6. What kinds of blocks get ignored by the model? A post-hoc manual analysis would have been nice.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "s2Qlk3PtkEu", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "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 proposes Block Skim Transformer (BST), a variation of Transformer-based architectures for QA where a newly proposed module predicts the relevance of each 'block' of context text and masks less relevant blocks out. The authors show 1.6x speedup during inference across 5 QA datasets, with a small increase in accuracy.\n\nWhile I find the motivation and the high level idea of using Attention weights to decide block relevance interesting, there are too many gaps in both the idea and the writeup at this point.  Below I list a few of them:\n\n* The way the block relevance part of the model is trained, it should really only work for factoid or single-hop QA tasks. This is because a block is marked relevant only if it contains answer words. On the other hand, in multi-hop or mulit-fact QA tasks (which is where the focus on the QA community currently is), this simple relevance criteria just doesn't work. E.g., in a question like, \"Where was the 44th president of the USA born?\", any information about WHO the 44th president is would not be marked as relevant, yet it's critical to answer the question.  The same applies to so-called \"comparison\" questions in HotpotQA.\n\n* Many important details are missing, which makes it difficult to assess the setup and the value of the findings. E.g., HotpotQA's so-called distractor setting (which is the easier one) has 10 paragraphs as context. How are these even fed into your baseline and BST models, when you have a max token limit of 512? Similarly for TriviaQA, where there are many (six, I think) evidence documents from Wikipedia.\n\n* The technique, by design, needs all context in the bottom 1/3 or 1/4 layers of the Transformer architecture.  This means it cannot handle any longer context than the baseline; everything must fit into the 512 token limit, which is arguably the biggest bottleneck of most of the current Transformer architectures. In other words, the paper doesn't address this problem, which is more important than obtaining a 1.6x speedup.\n\n* Eqn (2) is confusing: what are a and b? Start and end of a block? If so, this computation is the sum of all self-attention within the block from a to b. The text mentions using the attention between the question and the block for assessing relevance. The precise equation for that cross-attention would be more valuable here.\n\n* Eqn (6) is confusing: I was expecting a quadratic term, like N^2 in the numerator and k^2 N^2 in the right hand side of the denominator. What is \"Layer\" in this equation?  Also, please check your simplification, it doesn't seem correct (note, e.g., that as k increases, your speedup expression on the right also increases, but it should decrease). The correct term seems something like L / (l + (L-l) k^2)  =  L / (k^2 L + (1-k^2) l).\n\n* The speedup of 1.6x, while valuable, doesn't seem like a make-or-break decision when deciding to use a Transformer model or not at inference time.  I was hoping for a larger gain when keeping only a fraction of the context.\n\n* What about training time -- is there any speed up there?  Do I understand correctly the answer is \"no\" because, I think, you keep all blocks during training time in case the relevance decision is incorrect?\n\n* Inference experiments are motivated by a \"mobile\" use case setting, where a CPU is used with batch size one.  Is it correct to assume there isn't any significant speedup on a GPU with a larger batch, e.g., when computing Dev/Test numbers for the 5 QA datasets considered?\n\nThe writing has many typos and small mistakes, and would benefit from through proof reading.\n\nA convincing clarification of the above concerns could change my mind, but at this point I don't think the paper is ready for publication.", "extra_scores": {}, "extra_text": {}}, {"review_id": "SKhDsgK1hof", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "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": "\nSummary: This paper presents the \"Block Skim Transformer\" for extractive question answering tasks. The key idea in this model is using a classifier, on the self-attention distributions of a particular layer, to classify whether a large spans of non-contiguous text (blocks) contain the answer. If a block is rejected by the classifier, it is excluded in subsequent layers of self-attention. During training, no blocks are thrown away and the classifier is applied to every layer to provide a regularization effect, which leads to small improvements in performance in 5 datasets. During inference, blocks are thrown away at a fixed layer. The reduction in sequence length leads to ~1.5x batch size 1 speed improvements.\n\n-----------------------------------\n\nStrengths of the Paper\n\n1. Lots of experiments, using 5 different datasets and many different pretrained LMs like BERT, RoBERTa and ALBERT. Small but consistent improvements in most settings. A detailed ablation study testing hyperparameters and different components of the model.\n\n3. Cool idea of dropping tokens entirely, this is relatively unexplored in the space of efficient transformer models.\n\n-----------------------------------\n\nWeaknesses of the Paper\n\n1. Training-Inference mismatch: It seems to me the proposed model has a strong mismatch during training and inference. During training, none of the blocks are dropped. During inference, a number of tokens are dropped at a particular layer (which is a inference time hyperparameter). It seems counter-intuitive to me that the subsequent layers of the model are able to successfully process a partial input, especially when this was not done during training. Also, this line is confusing \"we forward the skipped blocks directly to the last layer for the QA classifier.\" Isn't this another training-inference mismatch? During training the QA classifier got the final layer representations for every block, not the intermediate layer representation.\n\n2. Improvements are quite small: in terms of QA performance, the improvement is just 0.3 F1 on an average across datasets. Is this improvement statistically significant? Even in terms of speed, a 1.6x batch size 1 improvement seems relatively modest. Are the speed / accuracy improvements better on tasks requiring long input sequences? Intuitively, the Skim idea seems more well-suited to tasks having inputs much longer than 384 tokens, perhaps you could use some of the benchmarks in the efficient transformer papers?\n\n3. How does this method compare to the numerous efficient transformers that have been proposed (in terms of attention, or adaptive layers)? While I don't expect comparisons against all of them, atleast 1-2 comparisons should be done to ground the observations better. Modifying attention distributions to be sparse over tokens seems like a more general way of getting the same effect as the block skim idea, and this also avoids the training-inference mismatch. \n\n-----------------------------------\n\nOverall Recommendation\n\nThe idea of dropping tokens is interesting, but I'm not very convinced by the results. Performance improvements are modest and I'm really confused about the train-inference mismatch in the proposed architecture. More comparisons with a few alternative efficient transformer papers are needed to make this a stronger submission. Overall I'm leaning reject, but open to increasing my score slightly if I get more clarification regarding weakness #1.\n\n-----------------------------------\n\nMinor: I don't really understand Eq 2. What's the final vector that's used for logistic regression? How exactly is the aggregation done between the question and a block of text?\n\n----------------------------------\n\n**After Author Response**\n\nThank you for the detailed clarifications. I've raised my score to 5 since I found the random block skipping experiment quite interesting and surprising. However, weakness #2 and #3 are fairly important in my opinion and I don't think the response sufficiently addresses those weaknesses. I encourage the authors to show their method works on longer sequence datasets (where skimming intuitively makes sense) and compare against 1-2 efficient transformers using sparse attention, for better grounding the improvements in existing literature.", "extra_scores": {}, "extra_text": {}}, {"review_id": "My1AmLVY-Nf", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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 Block Skim Transformer (BST), a variant of Transformer that skips tokens for self-attention so that it can pay less attention to less import tokens and achieve speedup. Although previous work studied the skimming mechanism for RNNs, this is the first paper that uses such an idea for transformers.\n\nThe model is pretty straight-forward - a block skim module in each transformer layer has a token-wise classifier that identifies which token to skip. One thing to note is that this model is specific to question answering (as the title indicates), because the supervision for tokens to skip and not to skip are pre-determined by identifying whether the token is part of the answer or not.\n\nStrengths of the paper:\n1) The problem is well-motivated with extensive discussion of related work.\n2) The model is straight-forward and the description is easy to follow.\n3) The paper includes extensive experiments using six different types of base pretrained models and five QA datasets, and comprehensive ablations.\n4) The model achieves impressive speedup, e.g., without more than 0.5% performance degradation, the model achieves x1.4 speedup and x1.8 speedup with BERT base and BERT large, respectively.\n\nWeakness of the paper:\nThe most important concern I have is that the model requires pre-defined supervision for tokens in order to train BST module. For instance, this paper trains the model to skim all the tokens that are not the answer token (which is the reason that it is specific to question answering). This means that BST module is trained for an end task (finding the answer tokens), rather than being trained to identify less important tokens for the end task. This is also a main difference from previous work that incorporates skimming modules for RNN, as most of them treat blocks to skip as latent variables and hence are easy to generalize to any tasks. This makes the model significantly less novel, and this fact was not apparent in Section 1 and 2, before reading the actual description of the model.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a model for question answering where blocks of text can be skipped and only relevant blocks are further processed for extracting the answer span.  
 The reviewers mostly praised the general idea.  
 R3 raised concerns on generalizability of the presented approach.  
R4 raised several issues regarding presentation and clarity.  
R2 and R4 have concerns regarding execution and find some of the results unconvincing.  
While I don't necessarily share R2s concern on small improvements (improvements are still statistically significance), and despite the approach being very interesting, there are several issues that reviewers pointed out and wasn't resolved after discussions. 

Decision: Reject