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

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 stage-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-theart 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 ( OpenQA ) is a question answering task where the goal is to train a language model to produce an answer for a given question . In contrast to many question answering tasks , an OpenQA model is only provided with the question as its input without accompanying documents that contain the answer . One of the most promising approaches to OpenQA is based on augmenting the language model with an external knowledge source such as Wikipedia ( often referred to as the evidence documents ) . In this approach , the model consists of two core components ( Chen et al. , 2017 ) : ( i ) an information retrieval system to identify useful pieces of text from the knowledge source ( the retriever ) ; and ( ii ) a system to produce the answer given the retrieved documents and the question ( the reader ) . We can view such a model as a latent variable model , where the latent variables represent retrieved documents that are used to produce answers given questions ( Lee et al. , 2019 ) . End-to-end ( joint ) training of this model is challenging since we need to learn both to generate an answer given retrieved documents and what to retrieve . Previous work considers two potential solutions ( see Table 1 for a high-level summary ) . First , they adopt a stage-wise training , where the retriever is trained while freezing the reader and vice versa ( Karpukhin et al. , 2020 , Izacard and Grave , 2021b , a ) . Another 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . alternative is to constraint the reader to condition on each retrieved document individually1 ( Guu et al. , 2020 ) —sometimes with extra supervision for the latent variables in the form of the relevant document for a question ( Lewis et al. , 2020b ) . In this paper , we consider a retrieval-augmented question answering model that combines information from multiple documents when generating answers . Expectation-maximization ( Dempster et al. , 1977 ) offers a principled template for learning this class of latent variable models . We present EMDR2 : End-to-end training of Multi-Document Reader and Retriever ( §2 ) . EMDR2 iteratively uses feedback from the model itself as “ pseudo labels ” of the latent variables for optimizing the retriever and reader parameters . We use two estimates of the latent variables : ( i ) prior scores for updating the reader parameters and ( ii ) approximate posterior scores given all observed variables for the retriever parameters . We evaluate our proposed method by experimenting on three commonly used OpenQA datasets : Natural Questions , TriviaQA , and WebQuestions ( §3 ) . EMDR2 achieves new state-of-the-art results for models of comparable size on all datasets , outperforming recent approaches by 2-3 absolute exact match points . We also show that EMDR2 is robust to retriever initialization . It achieves high accuracy with unsupervised initialization , suggesting that supervised training of the retriever may not be an essential component of the training process as suggested in prior work ( Karpukhin et al. , 2020 ) . In summary , our contributions are as follows : ( i ) we present an end-to-end training method ( EMDR2 ) for retrieval-augmented question-answering systems ; ( ii ) we demonstrate that EMDR2 outperforms other existing approaches of comparable size without any kind of supervision on the latent variables ; ( iii ) we provide ablation studies for a better understanding of the contributions of different components of our proposed method ; and ( iv ) we release our code and checkpoints to facilitate future work and for reproducibility.2 EMDR2 is a framework that can be used to train retrieval-augmented text generation models for any task . We believe that our estimation technique in EMDR2 is also useful for learning similar latent variable models in other domains . 2 Model . Our proposed model EMDR2 consists of two components : ( i ) a neural retriever and ( ii ) a neural reader , which we train jointly in an end-to-end setting . Figure 1 shows an illustration of our model and training procedure . We discuss each component and our training objective in detail below . 1This makes marginalization over the latent variables easier since we only need to consider one document at a time rather than multiple documents at once . 2Our code is available at : https : //github.com/DevSinghSachan/emdr2 2.1 Neural Retriever : Dual Encoder . Let the collection of evidence documents be denoted by D = { d1 , . . . , dM } . Given a question q , the goal of the retriever module is to select a subset of documents Z ⇢ D to answer the question . We model the retriever as a dual-encoder network ( Bromley et al. , 1994 ) , where one encoder fq encodes the question and another fd encodes the evidence document ( to a vector ) . The retrieval score is defined as the dot product between the two resulting vectors : score ( q , di ; ) = fq ( q ; q ) > fd ( di ; d ) , ( 1 ) where = [ q , d ] denotes the retriever parameters . We select top-K documents for the question q from D based on the retrieval scores . We denote the set of retrieved documents by Z = { z1 , . . . , zK } . We use transformer encoders ( Vaswani et al. , 2017 ) as our fq and fd . Our transformer architecture is similar to BERT with 12 layers and 768 hidden size ( Devlin et al. , 2019 ) . We use the final representation of the first token ( i.e. , the standard [ CLS ] token from BERT ’ s tokenization ) as our question ( and similarly document ) embedding . Initializing fq and fd with BERT weights has been shown to lead to a poor retrieval accuracy ( Lee et al. , 2019 , Sachan et al. , 2021 ) . Therefore , we initialize the retriever with an unsupervised training procedure . We discuss our initialization technique in detail in §3.2 . 2.2 Neural Reader : Fusion-in-Decoder . The reader takes as input a question q and a set of retrieved documents ( to be read ) Z to generate an answer . Our reader is based on the Fusion-in-Decoder ( FiD ; Izacard and Grave , 2021b ) model , which is built on top of T5 ( Raffel et al. , 2020 ) . T5 is a pretrained sequence-to-sequence transformer that consists of an encoder ge and a decoder gd . In FiD , each retrieved document zk is first appended with its title ( tzk ) and the question : xk = [ CLS ] q [ SEP ] tzk [ SEP ] zk [ SEP ] , where [ CLS ] is used to indicate the start of a document and [ SEP ] is used as a separator for the different parts of the document as well as the final token . Each xk is then independently given as an input to the T5 encoder ge . The output representations corresponding to all of the retrieved documents are concatenated as : XZ = [ ge ( x1 ) ; . . . ; ge ( xK ) ] 2 R ( N⇥K ) ⇥H , where N is the number of tokens in each xk3 and H is the hidden size of the T5 encoder ge . In this work , we use the T5-base configuration with N = 512 and H = 768 . XZ is then given as an input to the T5 decoder gd . When generating an answer token , the decoder attends to both previously generated tokens ( i.e. , causal attention ) as well as the tokens encoded in XZ ( i.e. , cross attention ) . Since XZ contains information from multiple documents , the decoder has the ability to aggregate useful signals contained in multiple documents and jointly reason over them . We define the probability of the answer as : p ( a | q , Z ; ⇥ ) = TY t=1 p ( at | a < t , q , Z ; ⇥ ) , ( 2 ) where ⇥ denotes the reader parameters ( i.e. , T5 encoder and decoder ) and T is the number of answer tokens . We keep generating answer tokens until the decoder outputs a special EOS token or a pre-specified maximum answer length is reached . 2.3 End-to-End Training of Reader and Retriever . In contrast to previous work on generative question answering , we train both the reader and the retriever jointly in an end-to-end differentiable fashion . Denote our latent variable which represents a set of retrieved documents by Z and let Z be a possible value of Z . The marginal likelihood of an answer ( marginalizing over all the possible values of Z ) 3We truncate and pad as necessary such that every xk has the same length N . See §3.2 for details . E idence . is : p ( a | q ; ⇥ , ) = P Z=Z p ( a | q , Z ; ⇥ ) p ( Z | q ; ) . The goal of our training procedure is to find and ⇥ that would maximize the above objective . Exactly optimizing Eq . 3 is intractable as it is combinatorial in nature.4 For one particular value Z , the log-likelihood is simpler to compute : log p ( a | q , Z ; ⇥ ) p ( Z | q ; ) = log p ( a | q , Z ; ⇥ ) + log p ( Z | q ; ) . Expectation-maximization ( EM ) algorithm ( Dempster et al. , 1977 ) offers a solution to learning this latent variable model . In classical EM , we iteratively compute the posterior of Z given all observed variables and use it to update ⇥ and . We propose using two estimates of Z—Zreader and Zretriever—for updating the two components of the model ( reader parameters ⇥ and retriever parameters ) : log p ( a | q , Zreader ; ⇥ ) | { z } reader + log p ( Zretriever | q ; ) | { z } retriever . ( 3 ) In the first term , we set the value of the latent variable Z = Zreader based on the prior scores . In the second term , we seek to maximize an approximate posterior of Z = Zretriever . We discuss them in more detail below . Reader parameters ⇥ . For updating ⇥ ( the first term of Eq . 3 ) , we use the top-K documents with the highest individual scores ( as computed by Eq . 1 based on the current value of ) to construct Zreader . This is equivalent to relying on the prior p ( Z | q ; ) to estimate Zreader ( without using information from the answer a ) . We choose to use the prior to train reader parameters since the prior scores are also used at evaluation time to obtain the top-K documents . As a result , there is no mismatch between training and test computations when computing p ( a | q , Z ; ⇥ ) ( i.e. , Z that is used at test time is obtained in exactly the same way as Zreader = Ztop-K ) . Retriever parameters . For updating ( the second term of Eq . 3 ) , we propose to use the posterior estimate . In other words , we use additional information from a when evaluating Zretriever to train . Using the posterior allows our retriever to learn from richer training signals as opposed to relying only on the prior . We need to be able to compute p ( Zretriever | q , a ; ⇥ , ) to maximize the retriever parameters . However , computing this quantity is difficult since it is a probability of a set.5 Consider a set of K documents ( e.g. , Ztop-K ) , where zk denotes a document in the set . We approximate the maximization of the probability of the set by assuming that its probability is maximized if the sum of the probability of 4Contrast our objective with REALM ( Guu et al. , 2020 ) , where the reader only conditions on one retrieved document zk when generating an answer . In this case , the latent variable represents a document assignment instead of a set of retrieved documents . 5This is true whether we choose to use the posterior probability or the prior probability . each document in the set is maximized.6 With this approximation , we arrive at a simpler quantity : P K k=1 p ( zk | q , a ; ⇥ , ) . Note that using Bayes rule , we can rewrite:7 p ( zk | q , a ; ⇥ , ) / p ( a | q , zk ; ⇥ ) p ( zk | q ; ) . ( 4 ) The reader now only conditions on one document when computing the probability of an answer p ( a | q , zk ; ⇥ ) . This simpler reader uses the same parameters as the more sophisticated one ⇥ , but it only uses one document zk instead of a set of documents . To compute Eq . 4 , we first obtain K documents with the highest scores as computed by Eq . 1 based on the current value of . We compute the probability of document zk 2 Ztop-K as : p ( zk | q , Ztop-K ; ) ⇡ exp ( score ( q , zk ) /⌧ ; ) P K j=1 exp ( score ( q , zj ) /⌧ ; ) , ( 5 ) where ⌧ is a temperature hyperparameter and the approximation assumes that documents beyond the top-K contributes very small scores so we do not need to sum over all evidence documents M in the denominator ( which is in the order of tens of millions in our experiments ) . We then compute p ( a | q , zk ; ⇥ ) similarly to Eq . 2 . Overall training objective of EMDR2 . Combining the above derivations , our end-to-end training objective that we seek to maximize for a particular example becomes : L = log p ( a | q , Ztop-K ; ⇥ ) | { z } reader + log KX k=1 SG ( p ( a | q , zk ; ⇥ ) ) p ( zk | q , Ztop-K ; ) | { z } retriever , ( 6 ) where SG is the stop-gradient operator so that the reader parameters ⇥ are not updated to also perform well given a single document zk . The stop-gradient operator in the second term of EMDR2 has several benefits . First , the FiD reader is trained from the first term of the EMDR2 objective in which its likelihood is conditioned on all the retrieved documents , similar to how the reader is used at test time . Second , it also makes training faster since the backward pass which is computationally more expensive than the forward pass is not needed , which in turn reduces the usage of GPU RAM as intermediate activations need not be saved . Given a training example , we update ⇥ and by taking gradients of Eq . 6 with respect to ⇥ and in an end-to-end fashion . Intuitively , we train the reader to generate the correct answer given K highest scoring documents Ztop-K . For the retriever , we train it to select K documents which collectively has a high score of generating an answer ( since the sum over K is inside the log in the second term ) while taking into account feedback from the reader . Algorithm 1 summarizes our training algorithm . Algorithm 1 : End-to-end training of multi-document reader and retriever . Input : Model parameters ⇥ and , evidence documents D. while not converged do • Compute Ztop-K using the current retriever parameters . // E-step • Compute p ( a | q , zk ) for each zk using the current reader parameters ⇥ . // E-step • Update model parameters ⇥ and to maximize the log-likelihood in Eq . 6 . // M-step end

Summary: The paper proposes end-to-end multi-document retriever-reader training for open-domain question answering. The proposed model is based on a pipeline approach where the retriever based on a dual encoder retrieves a set of documents from a large text corpus and a reader processes them and generates the answer. Different from previous approaches which independently train the retriever and the reader (Karpukhin et al, Izacard & Gave, and others), the paper jointly trains them using the EM algorithm. It is also different from previous work with joint training (Lee et al, Guu et al, and others) in that the generation of the answer is conditioned on multiple documents, not just a single one. In particular, this difference makes joint training more difficult, so the paper proposes to use the EM algorithm to iteratively update the parameters of the retriever and the reader. The paper also came up with clever initialization tricks and asynchronous refreshes of the MIPS index, which leads to successful training of the proposed model. The proposed model is evaluated on widely-used open-domain QA benchmarks (NQ, TriviaQA, WebQ). The proposed model outperforms a set of other, recently-proposed models by a large margin, which is impressive given that the baselines are highly competitive.