[Paper Review] End-to-End Training of Multi-Document Reader and Retriever for Open-Domain Question Answering

Devendra Singh Sachan, Siva Reddy+3 more|arXiv (Cornell University)|Jun 9, 2021

Topic Modeling45 citations

TL;DR

The paper proposes EMDR2, an end-to-end differentiable framework that jointly trains a multi-document reader and a retriever for open-domain QA by modeling retrieved document sets as latent variables and using an EM-inspired training loop.

ABSTRACT

We present an end-to-end differentiable training method for retrieval-augmented open-domain question answering systems that combine information from multiple retrieved documents when generating answers. We model retrieval decisions as latent variables over sets of relevant documents. Since marginalizing over sets of retrieved documents is computationally hard, we approximate this using an expectation-maximization algorithm. We iteratively estimate the value of our latent variable (the set of relevant documents for a given question) and then use this estimate to update the retriever and reader parameters. We hypothesize that such end-to-end training allows training signals to flow to the reader and then to the retriever better than staged-wise training. This results in a retriever that is able to select more relevant documents for a question and a reader that is trained on more accurate documents to generate an answer. Experiments on three benchmark datasets demonstrate that our proposed method outperforms all existing approaches of comparable size by 2-3% absolute exact match points, achieving new state-of-the-art results. Our results also demonstrate the feasibility of learning to retrieve to improve answer generation without explicit supervision of retrieval decisions.

Motivation & Objective

• Motivate end-to-end training for retrieval-augmented open-domain QA that leverages multiple documents.
• Propose an EM-based framework to jointly train a neural retriever and neural reader.
• Show that end-to-end training yields better document selection and answer quality than stage-wise training.
• Demonstrate state-of-the-art results on Natural Questions, TriviaQA, and WebQuestions with a comparable model size.

Proposed method

• Use a dual-encoder retriever to select a top-K set of documents for a question.
• Adopt Fusion-in-Decoder (FiD) as the multi-document reader to generate answers from multiple retrieved documents.
• Formulate end-to-end training as latent-variable learning over document sets Z and apply an EM-like objective.
• Estimate reader-relevant latent sets Z_reader from prior scores and retriever-relevant latent sets Z_retriever from a posterior-like objective.
• Compute a practical objective that combines p(a|q,Z_topK;Θ) for the reader with a tractable approximation to Σ_k p(a|q,z_k;Θ) p(z_k|q;Φ) for the retriever.
• Leverage a stop-gradient operator to stabilize joint training and prevent reader signals from leaking into retriever updates during certain steps.

Experimental results

Research questions

• RQ1 Can end-to-end training of a multi-document retriever-reader system improve open-domain QA performance compared to staged training?
• RQ2 Does modeling the retrieved document set as a latent variable and using EM-inspired updates yield better document selection and answer quality?
• RQ3 How does the proposed EMDR2 framework perform across standard OpenQA benchmarks (NQ, TriviaQA, WebQuestions) with a comparable model size?
• RQ4 Is the approach robust to retriever initialization, including unsupervised pre-training?
• RQ5 What are the effects of the number of retrieved documents and initialization strategies on performance?

Key findings

• EMDR2 achieves new state-of-the-art exact-match results on Natural Questions, TriviaQA, and WebQuestions for models of comparable size.
• End-to-end training with latent-set retrieval signals yields better performance than FiD-based baselines with same architecture.
• EMDR2 is robust to unsupervised retriever initialization and still achieves strong results, reducing the reliance on supervised retriever pretraining.
• Using 50 retrieved documents with end-to-end training outperforms configurations using the same documents with non-end-to-end training, highlighting the benefit of joint optimization.
• Compared to FiD-KD, EMDR2 delivers higher EM scores and requires fewer documents and training cycles, offering efficiency and robustness advantages.
• Ablation analyses show that the retrieval objective and the end-to-end EM-like updates contribute to performance gains, while certain alternatives can diverge or underperform.