An Efficient Memory-Augmented Transformer for Knowledge-Intensive NLP Tasks (2022)

Abstract: Access to external knowledge is essential for many natural language processing tasks, such as question answering and dialogue. Existing methods often rely on a parametric model that stores knowledge in its parameters, or use a retrieval-augmented model that has access to an external knowledge source. Parametric and retrieval-augmented models have complementary strengths in terms of computational efficiency and predictive accuracy. To combine the strength of both approaches, we propose the Efficient Memory-Augmented Transformer (EMAT) – it encodes external knowledge into a key-value memory and exploits the fast maximum inner product search for memory querying. We also introduce pre-training tasks that allow EMAT to encode informative key-value representations, and to learn an implicit strategy to integrate multiple memory slots into the transformer. Experiments on various knowledge-intensive tasks such as question answering and dialogue datasets show that, simply augmenting parametric models (T5-base) using our method produces more accurate results (e.g., 25.8 → 44.3 EM on NQ) while retaining a high throughput (e.g., 1000 queries/s on NQ). Compared to retrieval-augmented models, EMAT runs substantially faster across the board and produces more accurate results on WoW and ELI5.

## introduction nlp tasks often require knowledge that is not explicitly provided with the input. for example, open-domain question answering (odqa) requires answering an open-domain question without given context passages @xcite , and likewise for open-domain dialogue @xcite . to handle such tasks, one key challenge is storing and accessing potentially large amounts of knowledge. one approach is a parametric method that trains a sequence-to-sequence generator to represent knowledge within model parameters. @xcite find that pre-trained language models (plms) learn a partial knowledge base in their parameters, but its coverage is limited. increasing model size can improve this issue @xcite @xcite ; however, larger language models require significant computational resources. retrieval-augmented models @xcite @xcite , on the other hand, retrieve relevant passages from an external knowledge source (e.g., wikipedia), and use the retrieved passages to inform generation. despite being more accurate, retrieval-augmented models are often significantly more costly computation-wise than their parametric counterparts, since they require retrieving, encoding, and integrating the external knowledge at inference time. to combine the strengths of both parametric and retrieval-augmented models, we propose efficient memory-augmented transformers (emats) -an extension to transformer-based models augmented with an efficient key-value memory module. emat first encodes the external knowledge source into key embeddings and value embeddings, to construct the key-value memory (section 3.1). we choose paq @xcite , a large collection of question-answering generated from wikipedia, as our knowledge source; and we encode the questions as keys and answers as values. the transformer model produces dense query vector, retrieves from the key-value memory (section 3.2), and integrates the returned dense key-value vectors at different encoder layers to enhance generation (section 3.3). different from previous approaches @xcite @xcite , our query representation is com-puted at an early transformer layer, whereas retrieved key and value embeddings are incorporated into the model at a later layer. this design only requires one forward pass through the transformer model, and allows memory retrieval to run concurrently with the transformer forward pass, and hence reduces the computational overhead (see fig. 1 for our architecture). with this architecture, it is important that the key-value memory accurately represent the knowledge source, and the transformer learns a strategy to incorporate the retrieved key-value representations into the model. therefore, we introduce pretraining tasks (section 4.1), which include autoencoding objectives to represent the questions and answers, and a question answering task to learn an implicit strategy to incorporate multiple key-value memory slots. our ablation study (section 6.1) shows that our pre-training tasks are crucial for the performance, and removing any of them could lead to more than 10pp drop in em score on odqa datasets. our contribution can be summarised as follows: i) we introduce emat, an efficient memory access module to augment the transformer architecture; ii) we exploit paq as our knowledge source, and propose pre-training objectives to encode qapairs as key-value memory and to learn integration strategy to incorporate multiple memory slots into transformers; iii) experimental results on various knowledge-intensive tasks show that our proposed method significantly outperforms vanilla transformer baselines, while retaining a similar inference speed. efficient oqda systems one simple approach to accelerate odqa is closed-book qa (cbqa) -a sequence-to-sequence model @xcite such as t5 @xcite or bart @xcite ) is fine-tuned on odqa data, by training it to produce the answer given the question. cbqa models are substantially faster than retrieve-and-read approaches. however, since they solely rely on their parameters to store factual knowledge, their capacity is limited by the model size, and hence they often produce less accurate results than retrieveand-read methods @xcite . another efficient approach is retrieving semantically similar questions from a large collection of qa pair and returning the corresponding answers. @xcite propose paq, a 65 million qa dataset that is constructed with the objective to cover the most probably-asked questions in wikipedia. repaq @xcite , a retrieval-based qa system built on paq, won the efficientqa competition @xcite in 2020, outperforming cbqa models by a large margin. in this work, we choose paq as our knowledge source, but different from repaq, we develop a generative model. our results show that emat outperforms repaq while matching its efficiency. @xcite show that the feed-forward network (ffn) layers in transformer-based language models behave similarly to like key-value memories, where keys capture input patterns, and values map to the output vocabulary. based on this finding, @xcite propose to extend the ffn layers by concatenating a dense representation of the corpus to the layer weights. @xcite introduce a neural module to access a fixed external memory, showing that it can lead to significant improvements on downstream generative dialogue modelling tasks. concurrently to our work, @xcite propose qamat, a method to augment transformer layers with a key-value memory network encoding question-answer pairs. qamat requires two inference steps through the encoder: one to retrieve memory values, and another for concatenating the retrieved values to the input. in contrast, our proposed method only requires a single inference steps, resulting in a significantly smaller computational footprint. empirically, we show that our method is ≈ 5 times faster than qamat, even when using fewer hardware resources. ## efficient memory-augmented transformer in this work we propose efficient memory-augmented transformer (emat), a model architecture that uses a key-value memory to store millions of dense question-answer representations to inform its predictions (see fig. 1 ). given an input sequence @xmath0, • • • , x |x| ), emat's encoder first produces a dense query q to retrieve from the memory m. the returned key-value representations corresponding to the retrieved k key-value pairs are @xmath1, • • • , z k ). finally, the decoder generates the target sequence ## analysis 6.1 ablation study we conduct ablation study on the pre-training steps and the results are shown in ## conclusions in this work, we propose the efficient memory-augmented transformer (emat) that combines the strength of parametric model and retrievalaugmented model. it encodes external knowledge into a key-value memory and exploits the fast mips search for memory querying. we introduce pretraining tasks to learn better key-value representations and integration of multiple memory slots into transformer. experiments on knowledge intensive tasks, including open-domain question answering, dialogue and long-form question answering, show both the accuracy and speediness of emat. in the future, we will seek to improve integrate more diverse knowledge into the memory and generalise our method to more downstream tasks.