An Efficient Memory-Augmented Transformer for Knowledge-Intensive NLP Tasks

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. Our code and datasets are available at https://github. com/uclnlp/EMAT.

• The paper introduces EMAT which integrates a key-value memory module to enhance retrieval and integration of external knowledge in Transformer models.
• It leverages innovative pre-training strategies to represent QA pairs, significantly improving performance on datasets like NaturalQuestions and TriviaQA.
• EMAT’s architecture offers improved scalability and efficiency while maintaining competitive accuracy compared to traditional and retrieval-augmented models.

An Evaluation of Efficient Memory-Augmented Transformer for Knowledge-Intensive NLP Tasks

The paper introduces the Efficient Memory-Augmented Transformer (EMAT), a novel approach designed to enhance NLP models in handling knowledge-intensive tasks. This approach aims to combine the strengths of traditional parametric models and retrieval-augmented models, addressing key challenges in NLP such as question answering (QA) without provided context and open-domain dialogue. EMAT enhances Transformer-based models by incorporating a key-value memory module that encodes external knowledge into dense representations, enabling effective retrieval and integration of relevant information during inference.

Key Contributions

The paper outlines several innovative components that contribute to the effectiveness of EMAT:

• Key-Value Memory Encoding: EMAT utilizes a key-value memory structure that encodes questions as keys and corresponding answers as values, thus facilitating efficient retrieval of external knowledge without exhaustive parameter encoding. This choice allows the model to generalize across various tasks by efficiently accessing pre-encoded knowledge.
• Pre-training Strategies: The paper highlights the importance of pre-training tasks, which include auto-encoding objectives to represent questions and answers, and a generation task that allows the Transformer to integrate retrieved memory slots effectively. The empirical results suggest that excluding these tasks significantly reduces model performance, thus underscoring their necessity for high-quality representations and integration.
• Integration Strategy: EMAT's design incorporates retrieved key-value embeddings at specific layers of the Transformer, enhancing the model's prediction capability without a prohibitive increase in computational resources. This integration is executed concurrently with model inference, minimizing the computational overhead.

Empirical Evaluation

The authors conduct extensive experimentation across multiple knowledge-intensive datasets including NaturalQuestions, TriviaQA, WebQuestions, Wizard-of-Wikipedia, and ELI5. The results demonstrate that EMAT significantly outperforms traditional parametric models such as T5 and BART, both in terms of predictive accuracy and efficiency. Notably, EMAT achieves substantial improvements on NaturalQuestions (up to 44.3 EM) and exhibits comparable performance to retrieval-augmented models like RAG while drastically reducing the computational complexity.

Implications and Future Directions

EMAT's ability to enhance accurate and efficient inference in NLP tasks has several implications:

• Scalability: By integrating key-value memory directly into the model architecture, EMAT offers a scalable solution that can be adapted to various knowledge domains without necessitating significant changes to the model's parameterization.
• Interpretability: The use of explicit memory allows for greater interpretability in model predictions, as it is possible to trace which pieces of stored knowledge contribute to specific decisions.
• Broader Applicability: While primarily evaluated on QA tasks, the methodology could extend to diverse NLP applications requiring extensive external knowledge, suggesting fertile ground for future research.
• System Resource Utilization: EMAT's design reduces reliance on extensive hardware, offering practical utility in settings with limited computational resources. This characteristic opens avenues for deployment in academic and industrial contexts where resource efficiency is paramount.

In summary, this paper offers a compelling approach to integrating external knowledge into Transformer-based models efficiently. The combination of thoughtful architecture design and innovative pre-training strategies positions EMAT as a potential standard for future research in scalable, knowledge-driven NLP systems. Continued exploration may consider expanding the knowledge sources and evaluating broader applicability to different NLP contexts, potentially refining EMAT's architecture for even greater efficiency and performance.