[Paper Review] Augmenting Pre-trained Language Models with QA-Memory for Open-Domain Question Answering

Wenhu Chen, Pat Verga+3 more|arXiv (Cornell University)|Apr 10, 2022

Topic Modeling4 citations

TL;DR

This paper proposes QAMAT, a QA-memory-augmented transformer that improves open-domain question answering by pre-training a model to retrieve and use question-answer pairs from a large memory store. It introduces a two-stage pre-training method that boosts performance on small datasets and enables multi-hop reasoning via iterative retrieval, outperforming prior QA-only models on complex benchmarks like Musique and HotpotQA.

ABSTRACT

Retrieval augmented language models have recently become the standard for knowledge intensive tasks. Rather than relying purely on latent semantics within the parameters of large neural models, these methods enlist a semi-parametric memory to encode an index of knowledge for the model to retrieve over. Most prior work has employed text passages as the unit of knowledge, which has high coverage at the cost of interpretability, controllability, and efficiency. The opposite properties arise in other methods which have instead relied on knowledge base (KB) facts. At the same time, more recent work has demonstrated the effectiveness of storing and retrieving from an index of Q-A pairs derived from text \citep{lewis2021paq}. This approach yields a high coverage knowledge representation that maintains KB-like properties due to its representations being more atomic units of information. In this work we push this line of research further by proposing a question-answer augmented encoder-decoder model and accompanying pretraining strategy. This yields an end-to-end system that not only outperforms prior QA retrieval methods on single-hop QA tasks but also enables compositional reasoning, as demonstrated by strong performance on two multi-hop QA datasets. Together, these methods improve the ability to interpret and control the model while narrowing the performance gap with passage retrieval systems.

Motivation & Objective

• Address the lack of large-scale supervised data for question-question retrieval, a key bottleneck in QA-memory-based models.
• Improve performance on low-resource QA benchmarks like WebQuestions by introducing a novel pre-training task for question retrieval.
• Enable compositional reasoning for multi-hop questions using a QA-memory-based system, overcoming the limitation of prior models that only handle single-hop queries.
• Develop a more efficient training pipeline than latent-retrieval methods, avoiding expensive asynchronous index updates.
• Demonstrate that QA pairs can serve as effective, scalable knowledge units comparable to KB triples but with broader coverage than text passages.

Proposed method

• Propose QAMAT, a T5-based encoder-decoder model augmented with a key-value memory storing question-answer pairs.
• Use a two-stage pre-training strategy: first pre-train the encoder on in-batch QA pairs to learn retrieval and interpretation, then fine-tune on a fixed global memory.
• Construct a large-scale pre-training corpus by generating questions from text passages using existing question generation methods.
• Train the model to mask answers and retrieve the correct QA pair from the memory, simulating retrieval during inference.
• Introduce QAMAT+, a multi-round retrieval extension that iteratively queries the QA memory to chain reasoning steps for complex questions.
• Use a fixed global memory index after freezing the encoder, enabling efficient inference and avoiding costly online index updates.

Experimental results

Research questions

• RQ1 Can a pre-training task for question retrieval significantly improve performance on small, low-resource QA benchmarks?
• RQ2 Can a QA-memory-augmented model perform multi-hop reasoning without relying on document retrieval or structured knowledge bases?
• RQ3 Does a two-stage pre-training pipeline outperform end-to-end latent retrieval in terms of efficiency and performance?
• RQ4 Can QA pairs serve as effective knowledge units that support both single-hop and compositional reasoning?
• RQ5 How does the number of retrieved QA pairs affect model performance and inference efficiency?

Key findings

• The two-stage pre-training strategy significantly improves downstream performance, achieving 44.5% EM on NQ and 53.2% on TriviaQA without fine-tuning.
• QAMAT+ achieves state-of-the-art performance on Musique (a challenging multi-hop QA dataset), surpassing both RePAQ and T5-CBQA.
• Even with K=20 retrievals, EM scores on NQ and TriviaQA saturate, indicating diminishing returns beyond a moderate number of retrieved pairs.
• Removing the in-batch pre-training stage leads to a significant drop in performance (e.g., 26.0% EM on WebQuestions), showing its critical role in generalization.
• The model achieves 40% EM on NQ and TriviaQA in zero-shot settings using only pre-training, demonstrating strong few-shot generalization.
• QAMAT+ outperforms single-hop QAMAT on multi-hop datasets, proving the effectiveness of iterative retrieval for compositional reasoning.