PaperlayerView exact revision on arXiv

Source-linked AI summary

End-to-End Training of Multi-Document Reader and Retriever for Open-Domain Question Answering

Devendra Singh Sachan, Siva Reddy, William Hamilton, Chris Dyer, Dani Yogatama

arXiv:2106.05346v2cs.CLcs.AIcs.IR

TL;DR

Open-domain question answering must learn both which documents to retrieve and how to generate answers from them. EMDR2 uses expectation-maximization for end-to-end multi-document reader–retriever training, achieving state-of-the-art results across three benchmarks without explicit retrieval supervision.

  • Problem

    Jointly learning document retrieval and answer generation is challenging because retrieved documents are latent variables and prior methods often train components stage-wise or restrict document conditioning.

  • Method

    EMDR2 iteratively estimates latent relevant-document sets with expectation-maximization, using prior scores to update the reader and approximate posterior scores to update the retriever.

  • Results

    EMDR2 achieves new state-of-the-art results on Natural Questions, TriviaQA, and WebQuestions, outperforming comparable approaches by 2-3 absolute exact match points.

  • Takeaways & Limitations

    EMDR2 demonstrates end-to-end learning of retrieval-augmented question answering without supervision of retrieval decisions and is proposed for related latent-variable models.

  • Takeaways & Limitations

    The approach is resource-heavy, requiring 16 GPUs, and its experiments focus only on open-domain question answering.

Abstract

from arXiv · show

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.

1 Introduction

Open-domain question answering retrieves evidence documents and uses a reader to answer questions, but jointly learning retrieval and generation is challenging. EMDR2 addresses this with end-to-end latent-variable training and reports state-of-the-art results without retrieval supervision.

  • Open-domain question answering: OpenQA asks a language model to answer questions without accompanying documents, using a retriever and reader over external evidence such as Wikipedia.The retriever identifies useful text, while the reader generates an answer from the retrieved documents and question.
  • Challenge: Joint training is difficult because the model must learn both what to retrieve and how to generate answers from retrieved documents.Retrieved documents are treated as latent variables in the model.
  • Prior approaches: Prior approaches either train retrievers and readers stage-wise or constrain readers to condition on documents individually, sometimes with retrieval supervision.These alternatives separate or restrict how retrieval and answer generation interact.
  • EMDR2: EMDR2 trains a multi-document reader and retriever end-to-end using expectation-maximization and model feedback as pseudo-labels for latent retrieval variables.Prior scores update the reader, while approximate posterior scores update the retriever.
  • Results: EMDR2 achieves new state-of-the-art results on Natural Questions, TriviaQA, and WebQuestions, outperforming comparable approaches by 2-3 absolute exact match points.The method also remains effective with unsupervised retriever initialization.
  • Implications: EMDR2 is presented as a general framework for training retrieval-augmented text generation models and for learning related latent-variable models in other domains.The paper also reports contributions including ablations and released code and checkpoints.

2 Model

EMDR2 jointly trains a neural retriever and multi-document reader end to end, using latent-variable estimation to connect document retrieval with answer generation. Its Fusion-in-Decoder reader aggregates information across retrieved documents, while expectation-maximization approximations make training tractable.

  • The retriever selects a top-K subset of evidence documents using dual-encoder dot-product scores.Question and document encoders produce vectors whose dot product defines retrieval scores.
  • The Fusion-in-Decoder reader independently encodes each retrieved document with its title and question, concatenates the representations, and lets the decoder reason across them.The reader uses a T5 encoder-decoder; the decoder attends to both generated tokens and the concatenated document representations.
  • EMDR2 jointly trains a neural retriever and neural reader in an end-to-end differentiable setting.
  • Because marginalizing over document sets is combinatorial and intractable, EMDR2 uses expectation-maximization with separate latent-variable estimates for reader and retriever updates.Classical EM alternates posterior estimation and parameter updates; EMDR2 uses Zreader and Zretriever for the two components.
  • The reader uses prior-based top-K documents, whereas the retriever uses an answer-informed posterior estimate to obtain richer training signals.The prior matches evaluation-time retrieval, while the posterior incorporates information from the answer.
  • A stop-gradient operator prevents the reader parameters from being updated by the single-document auxiliary term, while reducing backward-pass computation and GPU memory use.The reader is trained through the multi-document likelihood used at test time.

3 Experiments

Experiments on NQ, TriviaQA, and WebQ evaluate EMDR2’s training setup, retrieval choices, baselines, and alternative objectives. EMDR2 achieves strong results across datasets, including low-resource WebQ, while end-to-end training and initialization choices affect performance.

  • Datasets and setup: Experiments use Natural Questions, TriviaQA, and WebQuestions, with Wikipedia segments as evidence documents.The evidence collection contains 21,015,324 non-overlapping 100-word Wikipedia segments.
  • Datasets and setup: The experiments use base-sized retriever and reader configurations and retrieve 50 documents unless stated otherwise.Both components use 12 layers, 768-dimensional hidden states, and 12 attention heads.
  • Main results: Top-1 retrieval vastly improves performance over no retrieval, while increasing FiD retrieval to top-50 substantially improves performance further.This supports modeling retrieved evidence as a set rather than conditioning on only one document.
  • Main results: EMDR2 outperforms all reimplemented FiD variants by large margins across NQ, TriviaQA, and WebQ.The compared FiD variants use MSS, DPR, or MSS + DPR retriever initialization in a two-stage training procedure.
  • Main results: EMDR2 exceeds reported FiD-KD results by more than 2.5 points on NQ and TriviaQA, obtaining new state-of-the-art results on those benchmarks.FiD-KD uses multiple training stages and knowledge distillation with inter-attention scores.
  • Efficiency and robustness: EMDR2 uses 50 evidence documents and one end-to-end training cycle, whereas FiD-KD uses 100 documents and multiple retriever-reader training cycles.EMDR2 also achieves state-of-the-art results with unsupervised retriever initialization.
  • Low-resource setting: On low-resource WebQ, EMDR2 improves over RAG by 3.5 points without supervised transfer from NQ.This result demonstrates applicability when only limited training examples are available.
  • Retriever initialization: Unsupervised MSS initialization improves retriever recall after EMDR2 training, while DPR initialization may not be essential for strong OpenQA performance.On NQ, MSS recall improves by 20% after training; similar trends are observed on TriviaQA and WebQ.

4 Related Work

Related work covers neural retrievers, neural readers, and retrieval-augmented applications beyond question answering.

  • Neural retrievers: Neural retrievers include dual encoders with one document embedding and multivector encoders with multiple embeddings.Multivector encoders can be computationally expensive for large-scale retrieval.
  • Neural readers: Neural readers are either extractive, producing an answer span from retrieved documents, or generative, conditioning answer generation on those documents.
  • Other application areas: Retrieval-augmented methods have also been applied to left-to-right language modeling, machine translation, and dialog.
  • The paper provides detailed comparisons with previous work in Appendices C and D.

5 Discussion

The discussion summarizes EMDR2’s end-to-end method and state-of-the-art results, while identifying computational, scope, resource, and societal limitations.

  • Summary of contributions: EMDR2 is an end-to-end training method for retrieval-augmented question answering systems.
  • Summary of contributions: EMDR2 uses expectation-maximization to derive its training objective and achieves state-of-the-art performance on three benchmark OpenQA datasets.
  • Technical limitations: Maintaining uncompressed evidence documents and searching for relevant documents can require substantial compute and memory.
  • Technical limitations: The experiments focus on open-domain question answering, leaving performance on other text generation models for future investigation.
  • Potential negative societal impacts: EMDR2 may exhibit large-language-model biases, lacks explicit fairness calibration, and could generate fake answers if evidence collections are modified by attackers.

Funding Statement

The paper acknowledges support from the Canada CIFAR AI Chair held by Prof. William Hamilton.

  • DSS was supported by the Canada CIFAR AI Chair held by Prof. William Hamilton.
  • The funding support is attributed to the Canada CIFAR AI Chair.
  • Prof. William Hamilton is identified as the holder of the cited Canada CIFAR AI Chair.

Checklist

The checklist records disclosures about claims, ethics, reproducibility, assets, datasets, and experimental reporting.

  • Claims and scope: The authors state that their main claims accurately reflect the paper’s contributions and scope.
  • Ethics and societal impacts: The paper reports its limitations and potential negative societal impacts in Section 5.
  • Reproducibility: The authors state that code, data, instructions, and training details are included in the supplementary material or relevant sections.
  • Experimental reporting: The experiments do not report error bars because repeated runs with different seeds were considered infeasible for compute reasons.
  • Assets and licensing: The paper uses open-source data and framework assets and describes license information when applicable.
  • Data and participants: The checklist marks consent, personally identifiable information, offensive content, crowdsourcing, human subjects, and participant compensation as not applicable.
  • Dataset statistics: Table 5 distinguishes the training set used for end-to-end QA training from the filtered set used for supervised retriever training.

A Dataset Details

The experiments use established OpenQA datasets and their standard splits, with documented preprocessing and openly available sources. Dataset sizes are reported for training, development, and test sets.

  • Validation examples are randomly selected from approximately 10% of each training set, using dataset splits from Lee et al. (2019).
  • The experiments use Natural Questions, TriviaQA, and WebQuestions, with dataset sizes provided for training, development, and test sets.
  • TriviaQA preprocessing uses human-annotated training answers and removes questions whose answers exceed 5 words.This filtering removes 2,362 training examples.
  • The datasets and preprocessed evidence resources are open-source and linked to their original or community-provided repositories.

B Additional Training Details

The paper documents implementation and training settings for its BERT, ICT, T5, and MSS components, including hyperparameter tables and substantial computational requirements.

  • BERT and ICT implementations are derived from Megatron-LM, with ICT initialized from BERT weights and trained on Wikipedia paragraphs truncated to 256 tokens.
  • T5 is implemented from Megatron-LM and trained for the original paper’s number of steps and batch size, using BERT lowercase tokenization.
  • The paper provides hyperparameter tables for BERT, ICT, T5, MSS, and dataset-specific fine-tuning.
  • MSS pre-training initializes the retriever from ICT weights and the reader from T5 weights before masking named entities in evidence sentences.
  • The supervised-training temperature parameter τ is set to the square root of the hidden size, following Sachan et al. (2021).
  • Experiments use 96 CPUs, 1.3TB memory, and 16 A100 GPUs; unsupervised MSS pre-training takes roughly one week.

C Unsupervised Pre-training and Comparisons with REALM

EMDR2 incorporates masked salient spans pre-training and asynchronous evidence-embedding updates, adapting techniques introduced in REALM for its own training procedure.

  • EMDR2 uses masked salient spans pre-training to improve retriever initialization and asynchronous evidence embedding updates during training.
  • These techniques were originally introduced in the REALM paper and are applied to EMDR2 with both similarities and differences.
  • The paper treats MSS pre-training and asynchronous evidence updates as components of its EMDR2 training procedure.

C.1 ICT and MSS Pre-training

ICT and MSS provide unsupervised retriever bootstrapping, while EMDR2 differs from prior systems through dense end-to-end retrieval and multi-document selection. Its objective yields stronger reported gains than Joint Top-K.

  • ICT and MSS Pre-training: ICT and MSS are unsupervised techniques used to bootstrap the retriever with good initial recall.
  • ICT and MSS Pre-training: ICT initialization produces higher Recall@5 than reported for REALM, which the authors associate with 768-dimensional evidence embeddings.
  • ICT and MSS Pre-training: After MSS training, EMDR2 achieves similar Recall@5 to REALM despite using a smaller batch size and fewer training steps.
  • ICT and MSS Pre-training: On WebQ, the MSS-pre-trained reader gains more than 1 EM point over the T5 reader.
  • ICT and MSS Pre-training: EMDR2 stores 21M evidence documents with 768-dimensional embeddings, requiring approximately 30 GB for FP16 evidence embeddings.
  • Comparisons with Prior Models: Unlike TF-IDF- and BM25-based prior systems, EMDR2 uses a trainable dense retriever and selects a set of documents for a multi-document reader.
  • Comparisons with Prior Models: On NQ and TriviaQA, EMDR2 improves over the FiD baseline by 4.3 and 6.4 EM points, while Joint Top-K gains 1 point and none, respectively.

E Qualitative Analysis

The qualitative analysis examines retrieval outputs and training-loss trajectories under different retriever initializations. Fine-tuning the MSS-pretrained retriever on NQ improves retrieval relevance, while initialization has only a marginal effect on answer-generation performance.

  • After MSS pre-training, top-1 retrieved documents relate to the question but are not relevant enough to answer it.The comparison uses representative NQ retrieval outputs.
  • NQ fine-tuning with EMDR2 improves retrieval accuracy, producing top-1 documents that are much more relevant for answering questions.
  • Comparing retriever initializations: MSS pre-training yields lower retriever accuracy than initialization with MSS followed by supervised DPR training.
  • Comparing retriever initializations: Retriever initialization has a marginal effect on answer-generation performance across NQ, TriviaQA, and WebQ.Reader training-loss plots compare MSS initialization with MSS + DPR initialization.
  • Visualizing reader and retriever losses: Figure 3 visualizes reader and retriever training-loss trajectories for a model initialized with MSS pre-training.
  • Table 10 marks answer-containing documents in blue and document probabilities from Eq. 5 in orange.
Loading 2106.05346v2…