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: The paper tackles the problem of open-ended knowledge-grounded natural language generation, in the context of free-form QA or knowledge-grounded dialogue, where models must ground their generations on passages relevant to the input context. Specifically, the authors explore improving the retrieval component of retrieval-augmented systems by utilizing posterior signal from the label. A “guide retriever” learns the relevant passage to retrieve by including the target output in its input context, and then, using an ELBo loss, provides this signal to the normal, in-system retriever. The authors find that their model improves three-fold over baselines: 1) it retrieves relevant passages more frequently; 2) its generations are more grounded in retrieved passages; and 3) its generations are closer to human generations.