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: This paper makes end-to-end training work for multi-document readers (e.g., FiD style readers). The key idea is an EM algorithm to update the reader and retriever given the top-k selection of previous iteration models. The proposed approach shows good performance on multiple open-domain QA benchmarks.