Hindsight: Posterior-guided training of retrievers for improved open-ended generation

1 INTRODUCTION . In knowledge-intensive NLP tasks , models must use open-domain knowledge to answer questions ( Kwiatkowski et al. , 2019 ; Joshi et al. , 2017 ) , fact-check claims ( Thorne et al. , 2018 ) , engage in informative conversations ( Dinan et al. , 2019 ) , or otherwise facilitate information access through natural language ( Nguyen et al. , 2016 ) . To leverage the information present in human-readable corpora ( e.g. , Wikipedia ) , many recent models for open-domain question answering ( Lee et al. , 2019 ; Khattab et al. , 2021 ) are retrieval-augmented : they extract passages from the corpus using a learned retriever and process it with a task-specific reader . If the relevant passage is known , the retriever is supervised to find it ; otherwise , it is weakly supervised with passages that contain the answer string ( e.g. , a name or a date ) . In recent work ( e.g. , Lewis et al. , 2020 ) , for end-to-end tasks with arbitrary text output , where relevant passages are not known , a jointly-trained generative sequence-to-sequence model provides the relevance signal . Current methods work well for short-answer QA-like problems , where the query itself has high overlap with the relevant passage and there is a unique answer , for example factoid Open-domain QA such as Natural Questions ( Kwiatkowski et al. , 2019 ) , HotPotQA ( Yang et al. , 2018 ) , fact-checking ( Thorne et al. , 2018 ) , and slot-filling ( Levy et al. , 2017 ) . However , for open-ended generation tasks like free-form QA ( MS-Marco NLGen ) , informative dialogue ( Dinan et al. , 2019 ) , or Wikipedia abstract generation ( Liu * et al. , 2018 ) , many responses can be equally acceptable and yet the data source ( e.g. , human-written answers , real-world conversations , Wikipedia abstracts ) supplies only one plausible output from a space that includes many . Figure 1 illustrates this with a conversational context x and multiple relevant passages z : while only one of them zgold is useful for generating the gold-standard answer y , other passages could have led to other coherent responses . This setting with one input and many possible unobserved outputs ( one-to-many ) is challenging , as it creates a distinction between the numerous passages that are relevant to the input ( which we dub contextrelevant passages ) and the few passages that pertain to the observed target output ( which we dub label-relevant passages ) . Had we known zgold corresponding to the target output , we could have supervised the retriever with zgold and trained the generator conditioned on zgold , but we don ’ t ! Current methods ( Lewis et al. , 2020 ) attempt to use the generator ’ s probability distribution Pθ ( y|x , z ) as a proxy for label relevance and train the retriever Pη ( z|x ) by marginalizing p ( y|x ) over retrieved documents z : P ( y|x ) = ∑ z∈top-k ( Pη ( .|x ) ) Pη ( z|x ) Pθ ( y|x , z ) . We find empirically that for one-to-many tasks , this objective is suboptimal in three ways : the generator is less grounded in the retrieved passages ( Figure 3 ) , the retriever performance saturates at low recall ( Figure 3 ) , and at training time the top-k retrieved passages miss out on many label-relevant passages weakening the supervision ( Table 1 ) . We propose explicitly estimating the label-posterior distribution Q ( z|x , y ) in the form of a guideretriever that has access to the target output , and with this “ hindsight ” can capture label-relevance precisely . We jointly optimize the retriever , posterior-guide , and generator using the evidence lower bound ( ELBo ) : Ezi∼Q ( .|x , y ) [ logPθ ( y|x , z ) ] − DKL ( Q|Pη ) . In theory , posterior-guided training is better for joint-training on open-ended tasks on all three counts : ( 1 ) The generator gets conditioned on passages weighted directly by their label-relevance from the label-posterior distribution , reducing the tendency of the generator to ignore the retrieved passage ; ( 2 ) The retriever is trained with reverse-KL divergence , whose mode-seeking nature encourages the retriever ’ s probability distribution to match some modes with the guide ( label-relevant passages ) , without explicitly penalizing other modes ( other context-relevant passages ) ; ( 3 ) As the label-posterior distribution Q ( z|x , y ) is modeled by a full-fledged retriever , it can retrieve label-relevant passages from the entire collection , which is a generalization of weak supervision approaches that use span overlap for short-answer QA . Our main contribution in this paper is the HINDSIGHT training system that : ( 1 ) uses a guide-retriever to provide a stronger learning signal for both the generator and the retriever , ( 2 ) is amenable to index-updates with iterative closed-set training ( Section 3 ) . We thoroughly evaluate the retrievalaugmented systems , going beyond end-to-end performance with a granular evaluation of the individual models ( retriever and generator ) at varying passage depths , important for evaluating one-to-many open-ended generation tasks ( Section 4 ) . Using HINDSIGHT on the Wizard of Wikipedia dataset of informative conversations our retriever achieves a 23 % relative improvement in success @ 10 ( its ability to have the label-relevant passage among the top-10 retrieved passages ) , our generator is more grounded with 19 % improvement in Novel-F1 overlap with the top-1 retrieved passage ( i.e. , its overlap with rare words from the retrieved passage that were not already in the input ) and 6.4 % relative improvement in Novel-F1 @ 1 overlap with the gold utterance ( the best matching generation when considering top-1 retrieved passage ) . We also validate that HINDSIGHT improves performance on the one-to-one generation task of free-form QA using the MS-MARCO NLGen dataset . 2 BACKGROUND . Open-domain Question Answering While Question-Answering datasets originally focused on extracting the answer from a given passage ( a.k.a . reading comprehension ) , recent work has made the task more challenging by not supplying a gold passage but instead expecting the model to con- duct open-domain QA directly over a large document collection . The first neural system for answering open-domain factoid questions was Dr. QA ( Chen et al. , 2017 ) , which used an off-the-shelf retriever ( e.g. , TF-IDF , BM25 ) to retrieve relevant passages and a separately trained reader to extract the answer span . The advent of efficient nearest-neighbour search algorithms ( Jegou et al. , 2010 ; Johnson et al. , 2017 ) opened up the possibility of encoding the query and the passage collection into a vector space and finding the relevant passage as nearest neighbors to the query embedding . Subsequent work trained neural retrievers in various ways such as : pretraining with the inverse cloze task then weakly supervising using span matches ( Lee et al. , 2019 ) , using gold passages with inbatch negatives ( Karpukhin et al. , 2020 ) , and retrieval-guided supervision with span-based positives ( Khattab et al. , 2021 ) . Open-ended Generation Natural language generation tasks involve generating a sequence of tokens ( maybe , word-pieces ) , often contextualized on some input ( another sequence of tokens for sequence-to-sequence tasks , an image for an image captioning task ) . However , they differ in how open-ended they can be . Some tasks such as factoid question-answering having a single correct short answer are less open-ended than free-form long answers . For example , in informative dialogue the speakers can lead the conversation in many different directions ( Dinan et al. , 2019 ) , also referred to as one-to-many generation . Thus it is more open ended , or , equivalently , it has a higher entropy than say machine translation Bojar et al . ( 2014 ) which has few correct translations that are very similar to each other . Many more generation tasks such as summarization ( Narayan et al. , 2018 ) and story generation ( Mostafazadeh et al. , 2016 ) lie on this spectrum . Retrieval for Language Modeling To improve the perplexity of a pre-trained language model , Khandelwal et al . ( 2020 ) retrieve similar contexts from the training set at each time-step and increase the likelihood of tokens that were predicted in similar contexts . Guu et al . ( 2020 ) instead pretrain a retrieval-augmented masked language model using salient-span masking and fine-tune it on downstream QA tasks . Under the paradigm of retrieval-augmented generation , for input x and output y , Lewis et al . ( 2020 ) retrieve top-k passages ( z ) from a corpus with a retriever and jointly train the generator ( Pθ ) and the retriever ( Pη ) by maximizing the likelihood of the output marginalized over the top-k documents , which we refer to as the MARGINALIZEDLOSS : P ( y|x ) = ∑ z∈top-k ( Pη ( .|x ) ) Pη ( z|x ) Pθ ( y|x , z ) ( 1 ) Here Pθ ( y|x , z ) has two roles : supervision , i.e . providing label-relevance by scoring label-relevant passages higher than other passages and grounding , i.e . maximizing the probability of the target output given context-relevant passages . For one-to-many datasets , where few context-relevant passages are label-relevant , the two roles are at odds with each other . By increasing label-relevance signal the generator reduces the ability to ground in passages that are context-relevant but not label-relevant . By increasing grounding in all context-relevant passages it reduces the label-relevance signal . Using label for direct supervision Zheng et al . ( 2020 ) use term-overlap with label as a heuristic to identify the gold-passage from a small set ( 5̃0 passages ) and train a reranker . In Zheng et al . ( 2021 ) they increase grounding by using the label to reweigh passage tokens and in Cai et al . ( 2019 ) they increase grounding by feeding a corrupted version of the label as a stand-in for label-relevant passage to the generator during training . In the next section , we build upon the end-to-end paradigm introduced by Lewis et al . ( 2020 ) with a theoretically motivated method to identify label-relevant passages from the entire corpus , and train a full-corpus retriever along with a grounded generator . 3 TRAINING WITH HINDSIGHT . To precisely identify label-relevant passages , we propose explicitly modeling the posterior distribution : Q ( z|x , y ) with a learned neural model . Unlike the retriever , the label-posterior model has access to the target output and in hindsight can differentiate label-relevant from other context-relevant passages . We learn the label-posterior jointly with the retriever and the generator by maximizing the evidence lower bound , ELBOLOSS , as given by the formula : logP ( y|x ) ≥ Ez∼Q ( .|x , y ) [ logPθ ( y|x , z ) ] −DKL ( Q|Pη ) ( 2 ) For an intuitive understanding we look at the two terms separately . The first term is an expectation of the generator ’ s log-likelihood Pθ over the label-posterior Q . This ensures that the generator need attend only to the label-relevant passages and therefore learns to “ trust ” the retrieved passages better . The second term is the KL divergence from the retriever to the label-posterior , also referred to as reverse KL divergence : DKL [ Q ( z|x , y ) | Pη ( z|x ) ] = ∑ z∼Q ( .|x , y ) Q ( z|x , y ) ( logQ ( z|x , y ) − Pη ( z|x ) ) This term is again weighted by Q ( z|x , y ) , which is the probabilistic equivalent of implication : high Q ( z|x , y ) implies high P ( z|x ) , i.e. , label-relevance implies context-relevance but not vice-versa . Posterior as a retriever While we can certainly model the label-posteriorQ ( z|x , y ) as a re-scorer that takes in documents as retrieved by the retriever Pη , we instead model it as a guide retriever that can retrieve label-relevant passages from the entire corpus . This is helpful because we can sample passages directly from the label-posterior distribution , and estimate the ELBOLOSS more accurately than using passages from Pη ( z|x ) . The guide retriever generalizes weak supervision approaches ( Lee et al. , 2019 ; Guu et al. , 2020 ) and relevance-guided supervision ( Khattab et al. , 2021 ) , to posterior-guided supervision with a learned posterior retriever rather than brittle heuristics based on word-overlap . Iterative closed-set training Prior work ( Guu et al. , 2020 ; Khattab et al. , 2021 ) has shown the utility of intermittently updating the passage index during training . To allow for such a workflow , we organize our training into rounds ( see Figure 2 ) . At the beginning of each round , in the outer loop we encode the passages and the queries with various retrievers and find the highest scoring r passages that we dub the closed-set . In the inner loop that runs for many epochs , we sample passages from the closed-set ( r = 100 for our experiments ) . This is fast because we are no longer retrieving from the entire corpus in the inner loop and also sufficient because the closed-set has a high recall . We update the retrievers ( both document and query encoders ) during the inner loop and use the latest model parameters for computing the loss functions . A round results in trained models that are then used for the next round . We find that 2 rounds are often sufficient , with decreasing marginal utility from the third round onward . Distributional repositioning before inference According to ELBOLOSS , the expectations are computed over z ∼ Q ( .|x , y ) . We can not compute the expectations exactly because of the large size of the passage corpus , so we approximate it by sampling k passages from the closed-set Qtop-r ( Q ( .|x , y ) ) . This leads to faster model training , especially at the beginning because passages from Q ( .|x , y ) provide better supervision . However , due to this sampling scheme , the models only ever get exposed to passages from the Q ( .|x , y ) distribution , which limits their ability to generalize over passages from Pη ( .|x ) during inference . To remedy this , we sample passages from an α-mixture of the two distributions : z ∼ { Pη ( .|x ) w.p . α Q ( .|x , y ) w.p . 1− α . In the initial rounds we set low values of α and increase it toward the end to reposition the passage distribution and better match with Pη ( .|x ) at test time . The retriever and the generator can be trained by sampling passages from different α-mixtures and we utilize this to avoid retriever overfitting ( with α = 1 ) while maintaining generator groundedness ( with α = 0.25 , 0.5 ) Training individual models to convergence A practical issue with joint training is that the retriever and generator train at different rates , reaching convergence at different times . The single term in MARGINALIZEDLOSS ( Eq . 1 ) does not indicate convergence of individual models leading to situations where one model starts to overfit while the other model still hasn ’ t converged . With ELBOLOSS there are two terms in Eq . 2 : the first connecting Q ( z|x , y ) and Pθ ( y|x , z ) , the second connecting Q ( z|x , y ) and Pη ( z|x ) . This allows us to fix the guide and train the retriever and generator independently for a differing number of gradient steps until each model reaches convergence .

Summary: In this paper authors describe an approach to use the responses/ answers to guide retrieval during training of document grounded response generator. By having the retriever being trained using the posterior (p(document|response,context)), the model could learn a better (supervised) retrieval network (based on ColBERT) which could be used to guide the training of the prior ( p(document|context)).  Documents are retrieved using the dialog context and the top-k documents are used for training the posterior as well as the prior. Since the expectation for ELBOLoss cannot be computed exactly (due to the large document set),  it is computed by sampling documents from the top-r retrieved documents -- from either the posterior or the prior, guided by a parameter called alpha which governs the sampling proportion. The idea of doing variational training using a posterior network isn't particularly novel but the authors have made it work for open-ended response generation using this approximation for computing the ELBOLoss. The networks use BART for language generation and CoLBERT for modeling the retrievers. Experiments have been presented using the Wizard of Wikipedia dataset and the MSMARCO NLGEN. Experiments show an improvement over the baseline model (RAG - referred to as MarginalizedLoss) in both retrieval (success@k, MRR) as well as response generation (text-F1 overlap between response and grounded document and textF1 overlap between response and ground-truth output response). Overall this is a well written easy to read paper