Title: FASTER AND JUST AS ACCURATE: A SIMPLE DECOM-POSITION FOR TRANSFORMER MODELS

Abstract: Large pre-trained Transformers such as BERT have been tremendously effective for many NLP tasks. However, inference in these large capacity models is prohibitively slow and expensive. Transformers are essentially a stack of selfattention layers which encode each input position using the entire input sequence as its context. However, we find that it may not be necessary to apply this expensive sequence-wide self-attention over at all layers. Based on this observation, we propose a decomposition to a pre-trained Transformer that allows the lower layers to process segments of the input independently enabling parallelism and caching. We show that the information loss due to this decomposition can be recovered in the upper layers with auxiliary supervision during fine-tuning. We evaluate decomposition with pre-trained BERT models on five different paired-input tasks in question answering, sentence similarity, and natural language inference. Results show that decomposition enables faster inference (up to 4x), significant memory reduction (up to 70%), while retaining most (up to 99%) of the original performance. We will release the code at <anonymized url>.

1. INTRODUCTION:

Inference in large Transformer-based NLP models such as BERT (Devlin et al., 2019) requires prohibitively high-levels of compute, making it expensive to support large volume processing in data centers, and almost infeasible to run on resource constrained mobile devices. These Transformer models create effective representations using self-attention, a mechanism that allows them to effectively account for wide textual contexts. However, applying self-attention over the entire input for all layers is computationally expensive. This raises a natural question: Is self-attention over the entire input necessary in all of the layers?

Previous studies (Tenney et al., 2019; Hao et al., 2019; Clark et al., 2019b) have shown that lower layers tend to capture syntactic phenomena that mostly depend on local contexts and that higher layers capture more semantic phenomena that are relevant to downstream tasks, which depend on longer global contexts. This suggests that considering only local context in lower layers of Transformer and considering full global context in upper layers can provide speedup at a very small cost in terms of effectiveness.

In this work we focus on paired-input NLP tasks such as reading comprehension, natural language inference and sentence pair similarity. These tasks provide a natural boundary for the locality of text (e.g., question vs. passage in QA). Because of this natural decomposition in two segments, we can compute representations for lower layers with only the local segment as the context and compute representations for upper layers with both segments as the context. This decomposition technique has multiple benefits: It allows for parallel processing of each segment, caching of segments that are available offline, and a significant reduction in runtime memory. Moreover, since the architecture remains largely same, the original pre-trained weights can be reused in the decomposed model. To compensate for the differences in the decomposed setting, we augment the fine-tuning loss on the target task with a distillation loss that minimizes the output-level as well as layer-level divergences.

We evaluate the decomposition idea using the BERT model on five different pairwise tasks. The decomposition achieves substantial speedup (2 to 4.3x) and reduction in memory (51.1% to 76.8%) for only small loss in effectiveness (0.2 to 1.8 points). Moreover, we find that with decomposition the larger BERT model can even run faster than the original smaller BERT model, while still being more accurate.

2. RELATED WORK:

Speeding up inference in a model requires reducing the amount of compute involved. There are two broad related directions of prior work:

(i) Compression techniques can be used to reduce model size through low rank approximation (Zhang et al., 2015; Kim et al., 2015; Tai et al., 2015; Chen et al., 2018) , and model weights pruning (Guo et al., 2016; Han et al., 2015) , which have been shown to help speedup inference in CNN and RNN based models. For Transformers, Michel et al. (2019) explore pruning the attention heads to gain inference speedup. This is an orthogonal approach that can be combined with our decomposition idea. However, for the paired-input tasks we consider, pruning heads only provides limited speedup. In more recent work Ma et al. (2019) propose approximating the quadratic attention computation with a tensor decomposition based multi-linear attention model. However, it is not clear how this multi-linear approximation can be applied to pre-trained Transformers like BERT.

(ii) Distillation techniques can be used to train smaller student networks to speedup inference. Tang et al. (2019) show that BERT can be used to guide designing smaller models (such as single-layer BiLSTM) for multiple tasks. But for the tasks we study, such very small models suffer a significant performance drop. For instance there is a 13% accuracy degration on MNLI task. Another closely related recent work is DistillBERT (Sanh et al., 2019) , which trains a smaller BERT model (half the size of BERT-base) that runs 1.5 times faster than the original BERT-base. However, the distilled model incurs a significant drop in accuracy. More importantly, this distillation model usually undergo expensive pre-training on the language modeling tasks before they can be fine-tuned for the downstream tasks.

In this work, we ask if can we speedup the inference of Transformer models without compressing or removing model parameters. Part of the massive success of pre-trained Transformer models for many NLP task is due to a large amount of parameters capacity to enable complex language representations. The decomposition we propose makes minimal changes retaining the overall capacity and structure of the original model but allows for faster inference by enabling parallel processing and caching of segments.

3. METHOD:

Transformers process the entire input sequence through multiple layers of self-attention, each of which involves an expensive transformation that is quadratic in the number of tokens in the input sequence. In the case of paired-input tasks with query and candidate texts (e.g. question and passage in QA, premise and hypothesis in NLI), the Transformers compute attention over a concatenation of both texts. This results in highly effective representations of the input pair, but this comes at a high-cost in terms of time complexity. If you want to reduce complexity, one natural question to ask is whether we can decompose the Transformer function over segments of the input, trading some representational power for efficiency. Consider the question-answering use case shown in Figure 1 , where the standard inputwide self-attention implicitly models question self-attention, the passage self-attention, and the two question and passage cross-attentions. The cross-attentions allow the question and passage to influence each other and produce effective representations for the target task. However, if we can process the two segments separately, we can improve efficiency in multiple ways. We get a basic reduction in compute because we no longer perform cross-attention, going from O(p + q) 2 to O(p 2 + q 2 ). More im-portantly we can get important gains through parallelism and caching. The question and passage self-attentions can be computed in parallel. Since the passage representation is no longer dependent on the question, they can be computed off-line and cached1 . The trade-off, however, is that we loose some representation effectiveness because we no longer use information from the other segment. We argue that this is a good trade-off from two viewpoints: Layer Variance (unit: 1e-4) (i) First, we can expect to achieve a good tradeoff in models with multiple layers, where the cross-attention is less important in lower layers when compared to the upper layers. Figure 2 supports this argument using the similarity of the representations of BERT layers for a single passage when computed using five different questions. The variance is smaller for lower layers and increases for higher layers. This is also in agreement with results on probing tasks which suggest that lower layers tend to model mostly local phenomena (e.g., POS, syntactic categories), while higher layers tend to model more semantic phenomena that are task dependent (e.g, entity co-reference) relying on wider contexts. We further posit that in these multi-layer models information loss in the lower layers can be potentially compensated for by the higher layers.

(ii) The approach requires minimal change to the overall Transformer architecture, which means we can reuse the pre-trained weights and expect that fine-tuning on the target task by itself will be sufficient to obtain high performance. The Transformer encoder has n layers (denoted L i for layer i), which transform this input sequentially:

3.1. DECOMPOSING:

X l+1 = L i (X l ).

For the details of the Transformer layer, we refer the reader to (Vaswani et al., 2017) . We denote the application of a stack of layers from layer i to layer j be denoted as L i:j . Full Transformer: The output representations of the full Transformer, A n and B n can be expressed as:

[A n ; B n ] = L 1:n ([A 0 ; B 0 ]) (1)

Decomposed Transformer: Figure 3 shows a schematic of our model. We decompose the computation of lower layers (up to layer k) by simply removing the cross-interactions between T a and T b representations. Here k is a hyper-parameter. The output representations of the decomposed Transformer, A n and B n can be expressed as:

[A n ; B n ] = L k+1:n ([L 1:k (A 0 ); L 1:k (B 0 )) (2) 3.2 SUPERVISING DECOMPOSED TRANSFORMER

The decomposed Transformer can be used in the same way as the original Transformer. Since the decomposed Transformer retains much of the original structure, we can initialize this model with the pre-trained weights of the original Transformer and fine-tune directly on downstream tasks. However, the decomposed Transformer looses some information in the representations of the lower layers. The upper layers can learn to compensate for this during fine-tuning. However, we can go further and use the original model behavior as an additional source of supervision.

Towards this end, we first initialize the parameters of decomposed Transformer with the parameters of a pre-trained full Transformer, and fine-tune it on the downstream tasks. To guide the learning of decomposed-Transformer further, we add auxiliary losses that makes predictions and the upper layer representations of decomposed Transformer closer to the predictions and corresponding layer representations of the full Transformer during fine-tuning.

Knowledge Distillation Loss: We use this loss to bring the prediction distribution of the decomposed Transformer closer to that of the full Transformer during fine-tuning. For this, we minimize the Kullback-Leibler divergence between decomposed Transformer prediction distribution P A and full Transformer prediction distribution P B :

L kd = D KL (P A P B )

Layerwise Representation Similarity Loss: We use this loss to bring the layer representations of decomposed Transformer closer to those of full Transformer during fine-tuning. For this, we minimize the euclidean distance between token representations of the upper layers of decomposed Transformer and the full Transformer. Let v i j be the token representations from i th layer (n layers) and j th token (m tokens) of the full Transformer. Likewise, let u i j be the token representations from i th layer (n layers) and j th token (m tokens) of the decomposed Transformer. Finally, let decomposed Transformer be decomposed up to layer k. We compute the Layerwise Representation Similarity loss for fine-tuning decomposed Transformer as following:

L lrs = n i=k m j=1 v i j -u i j 2

We add the Knowledge Distillation Loss and Layerwise Representation Similarity Loss along with the Task Specific Supervision Loss (L ts ) and learn their relative importance via hyper-parameter tuning:

EQUATION

We use the Bayesian Optimization (Močkus, 1975) to tune the γ, α and β instead of simple trialand-error or grid/random search. The Bayesian process is designed to minimize the number of steps required to find a combination of hyper-parameters that are close to the optimal one.

4. EXPERIMENTS:



4.1. DATASETS:

We use the pre-trained uncased BERT base and large 2 models on five different paired-input problems covering 3 QA tasks, 1 natural language inference task and 1 sentence-level semantic similarity task: For all 5 tasks, we use the standard splits provided with the datasets but in addition divide the original training data further to obtain a 10% split to use for tuning hyper-parameters (tune split), and use the original development split for reporting efficiency (FLOPs, memory usage) and effectiveness metrics (accuracy or F1 depending on the task).

SQuAD

4.2. IMPLEMENTATION DETAILS:

We implement all models in TensorFlow 1.14 (Abadi et al., 2015) based on the original BERT codebase (Devlin et al., 2019) . We perform all experiments on one TPU v3-8 node (8 cores, 128GB memory) with bfloat16 format enabled. We measure the FLOPs and memory consumption through the TensorFlow Profiler 3 . For decomposed Transformer models, we tune the hyperparameters for weighting different losses using bayesian optimizaiton libray (Nogueira, Fernando, 2019) with 50 iterations on the tune split (10% of the original training sets) and report the performance numbers on the original dev sets. The search range is [0.1, 2.0] for the 3 hyper-parameters.

For Decomposable BERT, we compute the representation for one of the input segments offline and cache it. For QA we cache the passages, for natural language inference, we cache the premise 4 and for question similarity we cache the first question 5 .

4.3. RESULTS:

Table 1 shows the main results comparing performance, inference speed and memory requirements of BERT-base and Decomposed BERT-base when using nine lower layers, and three upper layers (see Subsection 4.4 for the impact of the choice of upper/lower splits). We observe a substantial speedup and significant memory reduction in all the datasets, while retaining most of the original model's effectiveness (as much as 98.4% on SQuAD and 99.8% on QQP datasets). Efficiency improvements increase with the size of the text segment that can be cached.

Small Distilled or Large Decomposed? 3 : Inference latency (in seconds) on SQuAD datasets for BERT-base vs Decomp-BERT-base, as an average measured in batch mode. On the GPU and CPU we use a batch size 32 and on the phone (marked by *) we use a batch size of 1.

tive than using the smaller base model (+2.3 points) This shows that with decomposition, a large Transformer can run faster than a smaller one which is half its size, while also being more accurate.

We note that distilling a larger model into a smaller can yield better accuracy than training a smaller model from scratch. As far as we know, there are two related but not fully comparable results. (1) Tang et al. (2019) distill BERT to a small LSTM based model where they achieve 15x speedup but at a significant drop in accuracy of more than 13 points on MNLI. (2) Sanh et al. (2019) distill BERT to a smaller six layer Transformer, which can provide 1.6x speedup but gives >2 points accuracy drop on MNLI and >3 points F1 drop on SQuAD. A fair comparison requires more careful experimentation exploring different distillation sizes which requires repeating pre-training or data augmentation -an expensive proposition.

Device Results: To evaluate the impact on different devices, we deployed the models on three different machines (a GPU, CPU, and a mobile phone). Table 3 shows the average latency in answering a question measured on a subset of the SQuAD dataset. On all devices, we get more than three times speedup.

4.4. ABLATION STUDY:

Table 4 shows the contribution of auxiliary losses for fine-tuning Decomposed-BERT on SQuAD dataset. The drop in effectiveness when not using Layerwise Representation Similarity (LRS in table), and Knowlege Distillation (KD) losses shows the utility of auxiliary supervision.

Figure 4a and figure 4b shows how the effectiveness and inference speed of Decomposed-BERT changes as we change the separation layer. The plot shows that inference speed up roughly scales quadratically with respect to the separation layer number. The drop in effectiveness, on the other hand, is negligible for separating at lower layers (until layer 3 for the base model and until layer 13 for the large model) and increases slowly after that with a dramatic increase in the last layers closest to the output. The separation layer choice thus allows trading effectiveness for inference speed. 

5. ANALYSIS AND DISCUSSION:



5.1. UNDERSTANDING TRANSFORMER LOWER AND UPPER LAYERS:

The main difference between the original BERT and the decomposed BERT is the absence of cross attention in the lower layers. We analyze the differences between the representations of the two models across all layers. To this end, we randomly select 100 passages from SQuAD dev dataset as well as randomly selecting 5 different questions that already exist in the dataset associated with each passage. For each passage, we encode all 5 question-passage pair sequence using both the fine-tuned original BERT-base model and the decomposed model, and compute their distance of the vector representations at each layer. Figure 5 shows the averaged distances of both the question and passage at different layers.

Figure 5a indicates that the lower layer representations of the passage for both models remain similar but the upper layer representations differ significantly, supporting the argument that lower layers tend to capture more of local context than global contexts. In addition, the figure also shows that using the auxiliary supervision of upper layers has the desired effect of forcing the decomposed BERT model to produce representations that are closer to the original model.

Figure 5b shows the distance of question representations between original BERT and decomposed BERT, which also supports the findings from the passage representation. One minor difference is the smaller gap between using upper layer supervision for decomposed model and not using the supervision. We attribute this to the fact that question tokens are fewer than that of the passage, thus having less distance variations.

5.2. INFERENCE COST:

The decomposed model enables caching of text representations that can be computed offline. While a full-scale analysis of the detailed trade-offs in storage versus latency is beyond the scope of this Under review as a conference paper at ICLR 2020 Assuming a use case of evaluating one million question-passage pairs daily, we first compute the storage requirements of the representations of these passages. With the BERT-base representations we estimate this to be 226KB per passage and 226GB in total for 1 million passages. The cost of storing this data and the added compute costs and reading these passages at the current vendor rates amounts to a total of $61.7 dollars per month. To estimate inference cost, we use the compute times we obtain from our calculations and use current vendor rates for GPU workloads which amounts to $148.5 dollars to support the 1 million question-passage pair workload. The substantial reduction in cost is because the storage cost is many orders of magnitude cheaper than using GPUs. Details of these calculations are listed in the Appendix.

5.3. OTHER PRE-TRAINED TRANSFORMER MODELS:

More recently, two pre-trained Transformers XLNet (Yang et al., 2019) and RoBERTa (Liu et al., 2019) have outperformed BERT models on many NLP tasks. RoBERTa has the same architecture as BERT, it is reasonable to expect the decomposition idea to work accordingly. Although XLNet brings the segment-level recurrent into the Transformer self-attention layers that making it different from BERT, the decomposition technique is likely to be applicable due to being agnostic to the layer internals. We leave showing the effectiveness of decomposition of these two models to near future work.

6. CONCLUSION:

Transformers have improved the effectiveness of NLP tools by their ability to incorporate large contexts effectively in multiple layers. This however imposes a significant complexity cost. In this work, we showed that modeling such large contexts may not always be necessary and leverage this insight to build a decomposition of the Transformer model that provides substantial improvements in inference speed, memory reduction, while retaining most of the original model's accuracy. This decomposition model provides a simple yet strong starting point for efficient models as NLP moves towards increasingly larger models handling wider contexts.

It takes 226KB to store the vectors for 150 tokens 7 , and the total storage for 1 million sequences is 226GB. The Tesla V100 GPU allows a maximum batch size of 640 8 . We measure the t b = 4.6 for the original BERT-base model and t b = 1.4 for the decomposed BERT-base model. Then Cost original = 30 * 4.6 * 1, 000, 000/640 * 2.48/3600 = $148.5, and Cost decomp = 30 * 1.4 * 1, 000, 000/640 * 2.48/3600 + 30 * 1, 000, 000/10, 000 * 0.004 + 226 * 0.02 = $61.7.

7 vector dimension=768, bfloat16 format 8 >640 batch size will cause V100 GPU out of memory

:

There is a storage cost associated with caching. We show in Section 5.2 that storage cost is many orders of magnitude smaller compared to inference costs.

Class B operations on GCP

Figures and Tables:

Figure 1: Illustration of the attention matrix components for QA task.

Figure 2: Average variance of passage representations when paired with different questions at different layers.

A TRANSFORMER First, let's define formally the computation that a Transformer undertakes for a paired-task containing two segments of text, T a and T b . Let the token embedding representations of segment T a be A = [a 1 ; a 2 ; ...; a q ] and of T b be B = [b 1 ; b 2 ; ...; b p ]. The full input sequence X can be expressed by concatenating the token representations from segment T a and T b as X = [A; B].

Figure 3: Decomposing Transformers up to layer k, which enables encoding each segment independently from layer 1 to layer k. Auxiliary supervision of upper layer information from the original model further helps the decomposed model to compensate for information loss in the lower layers. KD is Knowledge Distillation loss and LRS is Layerwise Representation Similarity loss.

drop versus speedup on SQuAD for Decomposed BERT-large without auxiliary supervision.

Figure 4: F1 drop versus speedup of Decomposed BERT model (without auxiliary supervision) when separating at different layers.

Figure Representation distance of BERT vs decomp-BERT and distance BERT vs decomp-BERT w/o auxiliary loss/supervision

Table2compares performance, speed and memory of BERT-base, BERT-large and Decomposed BERT-large. Decomposed BERT-large is 1.6 times faster than the smaller BERT-base model. Decomposing the larger model turns out to be also more effec-(i) Performance of BERT-base vs Decomp-BERT-base, (ii) Performance drop, inference speedup and inference memory reduction of Decomp-BERT-base over BERT-base for 5 tasks. Decomp-BERT-base uses nine lower layers, and three upper layers with caching enabled. For SQuAD and RACE we also train with the auxiliary losses, and for the others we use the main supervision loss -the settings that give the best effectiveness during training. Note that the choice of the loss doesn't affect the efficiency metrics.

Table content: <table><tr><td/><td>Avg. Input</td><td>BERT</td><td>Decomp-</td><td colspan="2">Performance Drop</td><td colspan="2">Inference</td><td>Memory</td></tr><tr><td/><td>Tokens</td><td>base</td><td>BERT base</td><td colspan="2">(absolute | %age)</td><td colspan="2">Speedup</td><td>Reduction</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td>(times)</td><td>(%age)</td></tr><tr><td>SQuAD</td><td>320</td><td>88.5</td><td>87.1</td><td/><td>1.4 | 1.6</td><td/><td>3.2x</td><td>70.3</td></tr><tr><td>RACE</td><td>2048</td><td>66.3</td><td>64.5</td><td/><td>1.8 | 2.7</td><td/><td>3.4x</td><td>72.9</td></tr><tr><td>BoolQ</td><td>320</td><td>77.8</td><td>76.8</td><td/><td>1.0 | 1.3</td><td/><td>3.5x</td><td>72.0</td></tr><tr><td>MNLI</td><td>120</td><td>84.4</td><td>82.6</td><td/><td>1.8 | 2.1</td><td/><td>2.2x</td><td>56.4</td></tr><tr><td>QQP</td><td>100</td><td>90.5</td><td>90.3</td><td/><td>0.2 | 0.2</td><td/><td>2.0x</td><td>50.0</td></tr><tr><td/><td/><td colspan="3">Performance (Squad-F1)</td><td colspan="2">Speed (GFLOPs)</td><td>Memory (MB)</td></tr><tr><td colspan="2">BERT-large</td><td/><td>92.3</td><td/><td>204.1</td><td/><td>1549.6</td></tr><tr><td colspan="2">BERT-base</td><td/><td>88.5</td><td/><td>58.4</td><td/><td>584.2</td></tr><tr><td colspan="2">Decomp-BERT-large</td><td/><td>90.8</td><td/><td>47.7</td><td/><td>359.7</td></tr></table>

Performance, Inference Speed and Memory for different models on SQuAD.

Table content: <table><tr><td/><td>Tesla V100 GPU</td><td>Intel i9-7900X CPU</td><td>OnePlus 6 Phone</td></tr><tr><td>BERT-base</td><td>0.22</td><td>5.90</td><td>10.20*</td></tr><tr><td>Decomp-BERT-base</td><td>0.07</td><td>1.66</td><td>3.28*</td></tr><tr><td>Table</td><td/><td/><td/></tr></table>

Ablation analysis on SQuAD datasets for Decomp-BERT-base and Decomp-BERT-large models.

Table content: <table><tr><td/><td>Base Model</td><td>Large Model</td></tr><tr><td>BERT</td><td>88.5</td><td>92.3</td></tr><tr><td>Decomp-BERT</td><td>87.1</td><td>90.8</td></tr><tr><td>w/o LRS</td><td>86.2</td><td>88.9</td></tr><tr><td>w/o KD &amp; LRS</td><td>85.8</td><td>87.5</td></tr></table>

Peer review (iclr2020):
In this paper, a simple decomposition for transformer models is proposed to speed up the computation for a specific type of problems, where the input has paired, such as a question and an answer for QA tasks. The basic idea is to encode the question and the answer with self-attention alone, and then apply the co-attention on top of them, so that the complexity is reduced from O(p+q)^2 to O(p^2+q^2). Such a decomposed transformer is trained by distilling from a full transformer, where the knowledge distillation loss, layerwise similarity loss and task loss are adopted as the objective functions. Empirical studies show that the proposed method can be up to 3.5 times faster than the full transformer (BERT) model, while only losing a small amount of accuracy. 1. The paper lacks a large body of related and representative works in machine reading comprehension, where such a decomposition is widely used, such as BIDAF (Seo et al.), QANet (Yu et al.), R-Net (Wang et al.) among many others. If those works are properly acknowledged, the novelty in this paper will become slim. 2. The improvement from O(p+q)^2 to O(p^2+q^2) is negligible. On one hand, they are of the same order, so there is no theoretical contribution. On the other, as the author acknowledges, the method can only work on pair-input tasks like QA. However, most of the questions in QA are quite short, so the terms related to q can be ignored and the only bottleneck is p^2, making the difference very minor. 3. Since the method is essentially a distillation, the full transformer model is still needed for supervision. That is to say, it cannot work without the original large model. 4. There are many ways for distillation. For example, one can just use smaller hidden size or a shallower network, or their combinations. Those are straightforward baselines but they are not mentioned or compared with in the paper.