D e F ormer: Decomposing Pre-trained Transformers for Faster Question Answering (2020)

Abstract: Transformer-based QA models use input-wide self-attention – i.e. across both the question and the input passage – at all layers, causing them to be slow and memory-intensive. It turns out that we can get by without input-wide self-attention at all layers, especially in the lower layers. We introduce DeFormer, a decomposed transformer, which substitutes the full self-attention with question-wide and passage-wide self-attentions in the lower layers. This allows for question-independent processing of the input text representations, which in turn enables pre-computing passage representations reducing runtime compute drastically. Furthermore, because DeFormer is largely similar to the original model, we can initialize DeFormer with the pre-training weights of a standard transformer, and directly fine-tune on the target QA dataset. We show DeFormer versions of BERT and XLNet can be used to speed up QA by over 4.3x and with simple distillation-based losses they incur only a 1% drop in accuracy. We open source the code at https://github.com/StonyBrookNLP/deformer .

## introduction there is an increasing need to push question answering (qa) models in large volume web scale services @xcite and also to push them to resource constrained mobile devices for privacy and other performance reasons @xcite . stateof-the-art qa systems, like many other nlp applications, are built using large pre-trained transformers (e.g., bert @xcite , xlnet @xcite , roberta @xcite ). however, inference in these models requires prohibitively high-levels of runtime compute and memory making it expensive to support large volume deployments in data centers and infeasible to run on resource constrained mobile devices. our goal is to take pre-trained transformerbased models and modify them to enable faster inference for qa without having to repeat the pretraining. this is a critical requirement if we want to explore many points in the accuracy versus speed trade-off because pre-training is expensive. the main compute bottleneck in transformerbased models is the input-wide self-attention computation at each layer. in reading comprehension style qa, this amounts to computing self-attention over the question and the context text together. this helps the models create highly effective questiondependent context representations and vice-versa. of these, building representations of the context takes more time because it is typically much longer than the question. if the context can be processed independent of the question, then this expensive compute can be pushed offline saving significant runtime latency. can we process the context independent of the question, at least in some of the layers, without too much loss in effectiveness? there are two empirical observations that indicate that this is possible. first, previous studies have demonstrated that lower layers tend to focus on local phenomena such as syntactic aspects, while the higher layers focus on global (long distance) phenomena such as semantic aspects relevant for the target task @xcite @xcite . second, as we show later (see section 2), in a standard bert-based qa model, there is less variance in the lower layer representations of text when we vary the question. this means that in the lower layers information from the question is not as critical to form text representations. together, these suggest 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. based on these observations, we introduce de-former a simple decomposition of pre-trained transformer-based models, where lower layers in the decomposed model process the question and context text independently and the higher layers process them jointly (see figure 1 for a schematic illustration). suppose we allow k lower layers in a n-layer model to process the question and context text independently. deformer processes the context texts through k lower layers offline and caches the output from the k-th layer. during runtime the question is first processed through the k-layers of the model, and the text representation for the k-th layer is loaded from the cache. these two k-th layer representations are fed to the (k + 1)-th layer as input and further processing continues through the higher layers as in the original model. in addition to directly reducing the amount of runtime compute, this also reduces memory significantly as the intermediate text representations for the context are no longer held in memory. a key strength of this approach is that one can make any pre-trained transformer-based qa model faster by creating a corresponding deformer version that is directly fine-tuned on the target qa datasets without having to repeat the expensive pre-training. our empirical evaluation on multiple qa datasets show that with direct fine-tuning the decomposed model incurs only a small loss in accuracy compared to the full model. this loss in accuracy can be reduced further by learning from the original model. we want de-former to behave more like the original model. in particular, the upper layers of deformer should produce representations that capture the same kinds of information as the corresponding layers in the original model. we add two distillation-like auxiliary losses @xcite , which minimize the output-level and the layer-level divergences between the decomposed and original models. we evaluate deformer versions of two transformer-based models, bert and xlnet on three different qa tasks and two sentence-sentence paired-input tasks foot_0 . deformer achieves substantial speedup (2.7 to 4.3x) and reduction in memory (65.8% to 72.9%) for only small loss in effectiveness (0.6 to 1.8 points) for qa. moreover, we find that deformer version of bert-large is faster than the original version of the smaller bertbase model, while still being more accurate. ablations shows that the supervision strategies we introduce provide valuable accuracy improvements and further analysis illustrate that deformer provides good runtime vs accuracy trade-offs. ## decomposing transformers for faster inference the standard approach to using transformers for question answering is to compute the self-attention over both question and the input text (typically a passage). this yields highly effective representations of the input pair since often what information to extract from the text depends on the question and vice versa. if we want to reduce complexity, one natural question to ask is whether we can decompose the transformer function over each segment of the input, trading some representational power for gains in ability to push processing the text segment offline. the trade-off depends on how important it is to have attention from question tokens when forming text representations (and vice versa) in the lower layers. to assess this, we measured how the text representation changes when paired with different questions. in particular, we computed the average passage representation variance when paired with different questions. the variance is measured using cosine distance between the passage vectors and their centroid. as figure 2 shows that in the lower layers, the text representation does not change as much as it does in the upper layers, suggesting ignoring attention from question tokens in lower layers may not be a bad idea. 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 define the representation variance as the average cosine distance from the centroid to all representation vectors. in this figure, the variance is averaged for 100 paragraphs (each paired with 5 different questions) and normalized to [0, 1]. smaller variance in the lower layers indicates the passage representation depends less on the question, while higher variance in the upper layers shows the passage representation relies more on the interaction with the question. here we formally describe our approach for decomposing attention in the lower layers to allow question independent processing of the contexts. ## related work speeding up inference in a model requires reducing the amount of compute involved. there are two broad related directions of prior work: previous qa neural models like bidaf @xcite and many others contain decomposition as part of their neural architecture design. in contrast, the focus of our work is to show that large pre-trained transformer models can be decomposed at the fine-tuning stage to bring effectiveness of sota pre-trained transformers at much lower inference latency. 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. deformer applies to settings where the underlying model relies on input-wide self-attention layers. even with models that propose alternate ways to improve efficiency, as long as the models use inputwide self-attention, deformer can be applied as a complementary mechanism to further improve inference efficiency. we leave an evaluation of applying deformer on top of other recent efficiency optimized models for future work. ## 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. we 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. a key benefit of the model is that its architecture remains largely the same as the original model which allows us to avoid repeating pretraining and use the original model weights for finetuning. the distillation techniques further reduce the performance gap with respect to the original model. this decomposition model provides a simple yet strong starting point for efficient qa models as nlp moves towards increasingly larger models handling wider contexts.