Decoupled Transformer for Scalable Inference in Open-domain Question Answering (2021)

Abstract: Large transformer models, such as BERT, achieve state-of-the-art results in machine reading comprehension (MRC) for open-domain question answering (QA). However, transformers have a high computational cost for inference which makes them hard to apply to online QA systems for applications like voice assistants. To reduce computational cost and latency, we propose decoupling the transformer MRC model into input-component and cross-component. The decoupling allows for part of the representation computation to be performed offline and cached for online use. To retain the decoupled transformer accuracy, we devised a knowledge distillation objective from a standard transformer model. Moreover, we introduce learned representation compression layers which help reduce by four times the storage requirement for the cache. In experiments on the SQUAD 2.0 dataset, a decoupled transformer reduces the computational cost and latency of open-domain MRC by 30-40% with only 1.2 points worse F1-score compared to a standard transformer.

## introduction open-domain question answering (qa) aims to answer questions from a collection of text passages. it is an important and challenging task with application to several domains such as web search and voice assistants. the most popular architecture for open-domain qa is retriever-reader @xcite . given a question, the retriever uses an information retrieval (ir) system over a collection of passages to return top-k results that are most likely to contain an answer. then, the reader uses a machine reading comprehension (mrc) model on each of the top-k results to find an answer. in the end, the top-k mrc answers are ranked to produce a final answer. for both the retriever and the reader, large transformer models such as bert @xcite , roberta @xcite and electra @xcite achieve state-of-the-art results. a disadvantage of large transformer models is the high computational cost for inference which makes them hard to apply to online runtime systems, e.g. voice-assistants. transformers' computational cost comes from three major factors. firstly, the size of the feed-forward layers which project to an intermediate higher dimension and projects back to the original dimension. secondly, the multi-head selfattention has quadratic computational complexity in the sequence length. thirdly, the total number of layers. dense passage retrieval (dpr) @xcite ) is a retriever model which uses a transformer question encoder and transformer passage encoder to capture semantic similarity. the question and passage encoders are trained such that passages that are likely to contain an answer have a large embedding dot product with the question embedding. the embeddings for the passages are generated offline and indexed for efficient distributed knn search @xcite , and only the embedding for the question is generated at runtime. since questions are usually short, the retriever runtime inference computational cost is low. mrc reader models process the top-k passages returned by the retriever to get an answer. in transformer-based mrc models, each passage is encoded together with the question using a cls and separator characters [cls] question [sep] passage. the encoding is followed by a prediction head which determines the answer span. if there is no answer, the model result is a zerolength span on the cls token. the joint encoding of the document and the question produces rich interaction features but it increases the sequence length, and thus the computational coast. since the mrc model inference is executed at runtime on long sequences of a question and passages, mrc is the main computational bottleneck in retrieverreader qa. there have been several ideas to reduce the runtime inference of transformer models, such as precision reduction via quantization @xcite , knowledge distilling to a smaller architecture @xcite , and approximate multi-head attention for reducing the quadratic complexity @xcite . in this paper, we take an orthogonal approach, decoupling the transformer encoding for multiple inputs to improve efficiency, which can be combined with the aforementioned techniques. the motivation for the decoupled transformer is that in open-domain qa the passages are known in advance and part of the passage computation can be performed offline and stored. then, online at the runtime the question computation can be performed only once and combined with the stored state from passages with cross-attention. we use the decoupled transformer to reduce the computational cost of open-domain mrc by 30-40% with only 1.2 points worse than the f1-score on the squad 2.0 benchmark. our contributions are as follows: another model where the query and the passage are encoded independently using a transformer is colbert @xcite . the main modeling applications for colbert are retrieval and passage ranking. after encoding the query and the passage independently, late interactions are introduced using an efficient sum of maximum similarity computations. colbert representations are used for retrieval, so it combines the strengths of dpr and dc-bert. however, the efficient late interactions in colbert do not have enough representation power for complex tasks like mrc. ## decoupled transformer in the decoupled transformer, figure 1 , we split the transformer model m into two components. ## representation compression in open-domain qa the collection of passages are known in advance. so, with decoupled transformer, we run the input-component m input on each passage offline and store the passage representation in the index. for a large corpus, the representation storage can be a significant amount. in the case of qa over wikipedia, the storage requirement is around 3.4tb given around 32 million passages, averaging 150 tokens per passage, and 768 token dimensions for the bert-base model with 16-bit precision. to reduce the storage requirements for the passage representation of the decoupled transformer, we introduce a compression layer at the end of the input-component and a decompression layer at the start of the cross-component, see figure 2 . the compression layer is a linear projection from the original dimension to a compression dimension. the decompression layer is a linear projection from the compression dimension to the original dimension. these layers are similar to an autoencoder with a bottleneck. ## conclusion and future work we presented the decoupled transformer model for reducing runtime latency of mrc models in opendomain qa. the decoupling allows for part of the representation computation to be performed offline and cached for online use. to bridge the accuracy gap between a standard transformer and decoupled transformer, we devised knowledge distillation objectives for both model logits and features. moreover, we introduced a representation compression approach that allows for a four-times reduction in representation storage requirements for opendomain qa without significant loss of accuracy. we use the decoupled transformer to reduce the computational cost of open-domain mrc by 30-40% with only 1.2 points worse than the f1-score on the squad 2.0 benchmark. in the future, we are planning to extend the decoupled model with a dpr objective. the goal is for the input-component to also produce dpr-like embeddings suitable for similarity search. this way, we can have a single model that acts as both retrieval and reader.