Variational Learning for Unsupervised Knowledge Grounded Dialogs

1 Introduction

In this paper we focus our attention on the task of generating responses, grounded on information present in a large collection of external textual documents [Lewis et al., 2020]. In real-world scenarios, the exact document that one must access for generating the response is often unknown and one only has access to conversation logs and a document collection. Hence during training, given a dialog context, the primary challenge is first figuring out the correct document needed to generate the response, and then using that document for generating the actual response.

∗Contact Author

A straightforward baseline approach would be to use an out-of-the-box retriever (for instance, a tf-idf based retriever such as BM25 [Robertson et al., 1994] or a neural retriever such as DPR [Karpukhin et al., 2020]) for first retrieving the document and then using a retrieved document for generating the response. While this is fairly easy to implement, it cannot be trained in an end-to-end manner and thus, the retriever never improves as the model learns to generate responses.

To overcome this limitation, methods such as RAG [Lewis et al., 2020], model documents as latent variables and learn a distribution over these variables (Figure 1a). This distribution is referred to as the document prior. Specifically, the document-prior distribution is defined by querying a knowledge index [Johnson et al., 2017; Karpukhin et al., 2020] using the dialog context (history), and then converting the retrieval scores of the top-k 1 documents into a probability distribution. The response-likelihood can be defined using any neural language generator such as GPT2 [Radford et al., 2019]. It then performs a marginalisation of the latent variable over the retrieved documents to compute the approximate probability of the response, given the context. The negative log likelihood under this approximation forms the loss function to train.

However, one of the weaknesses of this approach is that, using the document-prior to query the index during training, ignores crucial information present in the ground-truth response which could have aided document retrieval. As a result, the response-likelihood network parameters may receive a weaker signal during training, which, in-turn, can cause models to try reducing their dependence on external knowledge by 'memorizing', especially if the correct document is rarely fetched by the retriever.

Variational Retrieval-Augmented Generation (VRAG): In this paper, we propose an approach that overcomes this limitation. We incorporate the ground truth response with the dialog context for retrieving the documents during training in a secondary retriever. This increases the chances of retrieving the correct document during training. The distribution over the documents defined by this retriever is referred to as the document posterior. The document posterior guides the training of the document prior while the documents sampled by the posterior are fed to the decoder for generating the response. Such a formalism emerges naturally in the variational setting, wherein the evidence lower bound (ELBO) is optimized instead of the maximum likelihood objective. Hence, we refer to the model as Variational Retrieval-Augmented Generation (VRAG).

1 typically k = 5-10.

One of the advantages of variational training is that it provides a low variance estimate of the objective (as compared to sampling from the document-prior distribution), for the same number of samples. Although this has been used in supervised settings ([Chen et al., 2020] and [Kim et al., 2020a]), we note that directly training under the variational objective may be prohibitively expensive in case of a very large 2 document collection (sampling an element from the posterior distribution, would require retrieval scores for each document in the index collection). In fact, related approaches for variational training therefore only use a small set of pre-retrieved documents [Lian et al., 2019; Li et al., 2020] to overcome this bottleneck. In particular, such approaches first use an out-ofthe box retriever to fetch a small set of documents (typically 5-10 documents) from the entire document collection. The methods then learn a prior as well as posterior distribution over the small set of pre-retrieved documents only by optimizing the variational objective.

A major weakness of this approach is that the out-of-thebox retriever does not benefit from training. As a result, if the recall of the out-of-the box retriever is low, that is, the correct document is not present in the pre-retrieved subset for most of the training data, the mapping from the document to the response learnt will be highly noisy and not very useful (we also demonstrate this experimentally in this paper).

Contributions: In this paper we describe our approach called VRAG or Variational Retrieval-Augmented Generation 3 which allows us to extend variational optimization to cases where documents are retrieved from large document collections. Instead of pre-retrieving a small set of documents to facilitate variational training, we retrieve documents from the entire collection but perform a summation over the top-k retrieved documents from the posterior distribution as well as the prior distribution to approximate the variational objective (Figure 1b). Top-k retrieval can be performed efficiently using an index for nearest neighbor search such as Faiss [Johnson et al., 2017], and we find that this simple trick performs significantly better than other approaches. We present experiments on three, publicly available, conversational QA datasets and we show that variational training helps build better knowledge grounded dialog systems. Our experiments show that not only does VRAG perform better on the end-task, it also learns a better retriever. To the best of our knowledge, we are the first to apply variational training for open-scale unsupervised knowledge grounded dialog systems. 4

2 Collections can have millions of documents

3 We provide the code and the supplementary material at https: //github.com/mayank31398/VRAG and https://arxiv.org/abs/2112. 00653 respectively.

4 Concurrently, [Paranjape et al., 2022] also use variational training with RAG to generate responses

2 Background

As is commonly done in dialog modeling tasks, we represent the collection of dialogs as a set of context (dialog history) and response pairs; T = {(x(i) , y(i) )}mi=1 where each context x(i) as well as its response y(i) is a sequence of tokens. Further let D = {dj}Nj=1 be a set of documents in the form of a large indexed document collection. We assume that each context-response pair requires exactly 1 document dj ∈ D (where 1 ≤ j ≤ N) to generate the corresponding response. Let z(i) denote a discrete variable which indicates the document (from the indexed collection) needed for training instance i i.e, dzi ∈ D. We can now model the joint likelihood of a response and document pair (y(i) , z(i) ) as p(y(i) , z(i) |x(i) ) = p(z(i) |x(i) )p(y(i) |x(i) , z(i) ).

In the absence of document-level supervision, the z(i) variables are unknown or 'latent'. Here we will be maximizing log p(y(i) |x(i) ) = ∑z p(z(i) |x(i) )p(y(i) |z(i) , x(i) ). However, since an explicit summation over the entire document collection can be computationally intractable, one needs to resort to a few approximation techniques. For ease of notation, we will drop the superscript (i) from now on.

Retrieval based approaches: Approaches such as RAG [Lewis et al., 2020] and REALM [Guu et al., 2020], maintain an index which allows one to retrieve the top-k documents with high prior probability. The objective is then approximated as a sum over these retrieved documents.

Specifically, the document-prior distribution p(z|x) is defined based on scores returned by the Dense Passage Retriever (DPR) [Karpukhin et al., 2020].

The top-k most relevant documents (Sp k ) for a query (dialog context) are retrieved from an index that allows efficient retrieval [Johnson et al., 2017] using MIPS search. We denote the approximate document-prior distribution, normalized over the set Sp k , as ˆ p(z|x). The overall objective for generating the response can then be written as log [∑z∈Sp k ˆ p(z|x)p(y|z, x)]. RAG suffers from a drawback that it does not use the information from responses in order to retrieve documents for a given training instance.

Variational techniques: An alternative approach is to maximize a variational lower bound on the objective. Here we need to define an Evidence Lower Bound (ELBO) on the likelihood as log p(y|x) ≥ Ez∼q(z|x,y) [log ( p(y,z|x) q(z|x,y) )]. This lower bound holds for any distribution q. Variational autoencoders [Kingma and Welling, 2013] define another network to model the distribution q. To train such networks, the ELBO is split as:

The first term is an expectation that can be estimated by sampling documents from the document-posterior q(z|x, y) distribution. The response-likelihood network is then run only using these sampled documents. One can either use re-parameterization trick (with Gumbel softmax distribution) [Jang et al., 2016; Maddison et al., 2016] or policy gradient method to back propagate through the sampling step. However, in order to sample a document, one would need to compute the entire distribution over the documents. This can be prohibitively expensive when using a large document collection. The second term (KL-divergence) is also computed as an explicit sum, given access to prior and posterior probability distributions, and is also intractable for large document collections.

In summary, variational training which uses the posterior distribution to retrieve documents while training, can help retrieve more relevant documents for training the responsegenerator. One of the trivial ways to extend these approaches (for large document corpus) is to identify the candidate knowledge documents for each training instance via an existing out-of-the-box retriever ([Li et al., 2020], [Lian et al., 2019]). One can then create prior and posterior distributions on the restricted set of documents. However this does not allow us to train the retriever and we also show in our experiments, that the trained distribution in such a setting does a poor job of generalizing as a retriever (Section 4).

In our approach we generalize the variational technique to open domain setting without fixing or pre-retrieving those candidate documents. In order to do so, we would first need to be able to compute the ELBO objective more efficiently (over the entire document corpus) as described in the next section.

3 Variational RAG (VRAG)

Variational training involves using both, the document-prior (p(z|x)) and document-posterior distributions (q(z|x, y)). We model each of these based on scores from a Dense Passage retriever (DPR) [Karpukhin et al., 2020] i.e.

where f and g are parameterized representations of documents (z) and dialog contexts (x). h(x, y) denotes the joint embedding of the context-response pair. These are created using neural models such as BERT [Devlin et al., 2018]. We use a single network to compute the document embeddings f(z) for both prior and posterior.

In order to efficiently compute expectation and KL divergence terms in the ELBO objective (Equation 1), we need to approximate the above distributions. To do so, we maintain an index on document embeddings. This allows us to retrieve the set of top-k documents under prior and posterior distributions (equations 2 and 3) using MIPS search [Johnson et al., 2017]. Note that since the query embeddings g and h are trainable, the retriever is trained over the epochs as well. We denote the sets of top documents under prior and posterior distributions by Sp k and Sqk respectively.

The overall cost is then computed using Evidence Lower Bound (Equation 1). Here for the first term, we normalize ˆ q over the set of top-k documents (denoted by Sqk ) returned by the index when queried using the posterior (i.e, the response and the dialog context). The first term is then approximated as ∑z∈Sqk ˆ q(z|x, y) log p(y|z, x). To approximate the KLdivergence we use the top-k knowledge instances retrieved by querying the index using the dialog-context (for prior) and the context-response pair (for posterior). We then take union of the two sets Sp k and Sqk to form the set SKL. We use this set (SKL) to approximate the KL-divergence. Thus, the KLdivergence in Equation 1 is given by:

where the approximate posterior (ˆ q) and prior (ˆ p) in this case are obtained by normalizing the retrieval scores on SKL. The intuition behind this approximation is that the documents with low posterior probability do not contribute much to the KL objective and hence, can safely be ignored. Similar to other variational models, VRAG is trained end-to-end.

3.1 Architecture

We now describe the neural networks used to model the prior and posterior distributions, and the response generator.

Prior Distribution Encoders: We need to create document and context representations defined by functions fand g, respectively for modelling the prior. We pass the input context through BERT model [Devlin et al., 2018] to create context representation. We use special markers to separate the turns. The embedding at final layer of [CLS] token is passed through a linear layer to create context representation. Similarly the document is passed through (separate) BERT Model to create a document representation at [CLS] token.

Posterior Distribution Encoders: Similar to the modeling of the prior distribution, we use another (separate) BERT to model the posterior. Here we create the input representation of the context-response pair (x, y), where a special marker is used to separate the context and response.

Response Likelihood: We use the standard sequence to sequence formulation for response likelihood, where log p(y|z, x) = ∑j log p(yj|y<j, z, x). Here, we use the GPT2 [Radford et al., 2019] model as our decoder with input sequence consisting of context and response.

3.2 Training Details

We train our network to maximize the ELBO objective (Equation 1). We initialize our document-prior (for both RAG and VRAG) and document-posterior (for VRAG) networks with the pretrained DPR-Multiset model 5 pre-trained using data from the Natural Questions [Kwiatkowski et al., 2019], TriviaQA [Joshi et al., 2017] etc.

During the training of the model, it can be difficult to rebuild the document index after every change to the document representation parameters in f, therefore similar to Lewis et al., the parameters in f are kept constant.

We used early stopping with patience = 5 on recall of the validation sets to prevent overfitting of models. The loss was optimized using AdamW Optimizer [Loshchilov and Hutter, 2017]. We also found it useful to continue training the response-likelihood for both RAG and VRAG after the joint training is complete. This is because while training, the decoder-likelihood function often lags behind prior (for RAG) and posterior (for VRAG).

5 This can be used to initialize a model in Hugging Face, see https://huggingface.co/transformers/model doc/dpr.html

(b) VRAG (Our approach)

Figure 1: (a) RAG does not use information from the responses to retrieve documents. (b) Our approach - VRAG which is trained end-to-end and uses the response to train a posterior distribution which guides the prior distribution

3.3 Response Generation

At test time we need to generate the response for a given context by first retrieving a document - thus, in both RAG and VRAG models we use the trained document-prior model to retrieve the top-k documents using the dialog context as query. We experiment with two different decoding strategies to generate the response:

1. Top Document Decoding: In this case the document with the highest prior probability is used to condition the generator. The response is then generated using beam search (beam width=3) on the trained GPT-2 response generation model; the most likely beam is taken to be the prediction of the model. We refer to this method as the 'top-1 decoding' in our experiments.
2. Top-k Documents Decoding: Here top (k=5) documents are retrieved from the prior distribution, say

z1, ..., z5. A beam search is then run to generate the top response from each of these say r1, ..., r5. We use the estimate of p(ri|x) ≈ p(z|x)p(ri|z, x). The most likely response under the estimated distribution is taken to be the response generated by the model 6 . We refer to this as 'top-5 decoding' in our experiments.

4 Experiments

Our experiments aim to answer the following questions: (1) Does Variational RAG (VRAG), which uses samples from the approximated document-posterior distribution, perform better than vanilla RAG? (2) Does the quality of generated responses improve by decoding using the top-k documents? (3) How do the trained document retriever modules for RAG

6 This is the same as "Fast Decoding" as defined in [Lewis et al., 2020].

and VRAG compare with each other? (4) Are the quality of samples returned by the document-posterior of VRAG better than document-prior of RAG as hypothesized? (5) How does VRAG compare with other approximations for variational training?

4.1 Datasets

OR-QuAC [Qu et al., 2020]: This dataset is a modified version of the QuAC [Choi et al., 2018] dataset. The dataset consists of dialog conversations, where each conversation is associated with the top-5 most relevant documents (retrieved using TF-IDF [Robertson et al., 1994] based BM25 ranking) from the QuAC dataset. To create an open-scale collection for our task, we index the set of all the documents available in the train, validation and test splits. In some cases of the test and validation set, the ground-truth document may be missing in the top-5 list associated with each conversation. In such cases, we obtain the ground-truth document from the original QuAC dataset and add it to the indexed collection.

DSTC9 [Kim et al., 2020b]: This dialog dataset was released as part of the DSTC9 challenge. The dataset comprises of dialog conversation turns in which the system: (i) needs to identify the turns in which to consult a collection for textual FAQs, (ii) retrieve the FAQ if required, (iii) and then generate the response based on the retrieved FAQ. The universe of knowledge documents in this dataset is the set of FAQs and each FAQ also includes the entity name because the same question can occur multiple times for different entities (eg: "Is parking available?"). The training dataset consists of conversations based on 4 different domains i.e hotels, restaurants, trains and taxis. The test dataset contains an additional domain, attractions, which is not found in the train and validation splits.

DoQA [Campos et al., 2020]: This dataset comprises of open-ended dialog conversations on different domains like cooking, travel and movies. Unlike, the OR-QuAC and DSTC9 datasets, most questions in this dataset are not factoid/specific questions, and are open-ended. We only use the cooking split for both training and testing.

We preprocess all the datasets by removing all the examples if the ground truth response is "CANNOTANSWER" (unanswerable). Each 'instance' refers to a context-response pair.

Question-Answering Task: The OR-QuAC dataset also contains non-contextual variations of questions at each dialog turn and we use them in a QA setting (referred to as ORQuAC-QA).

4.2 Baselines

Apart from our approach (VRAG), we also study the performance of RAG, as well as, a pipeline model which uses the pre-trained DPR-Multiset Retriever and a GPT2 based decoder which is fine-tuned to generate responses. We refer to this as the DPR + GPT2 baseline in our experiments. In addition, we also show the importance of using a trainable retriever for variational setting by comparing gainst a variational model where a fixed set of candidate documents are retrieved using a pre-trained DPR as a retriever (conditioned on context and response) during training [Lian et al., 2019;

Li et al., 2020]. This model retrieves from the entire document collection using the prior distribution during inference. We refer to this baseline as Variational Model with preretrieval ("pre-VM"). Thus, the distributions being trained do not change the set of candidate documents used as training progresses. All models use their respective document-prior distributions during testing.

4.3 EvaluationMetrics

For each of our experimental runs, we report the Mean Reciprocal Rank@5 (MRR@5), Recall@1 (R@1), Recall@5 (R@5) to evaluate prior's performance. We also report BLEU scores (both with top-1 and top-5 decoding) to assess the performance of generator. The BLEU-1 and BLEU-4 scores in our tables are denoted by B-1 and B-4 respectively. We also consider "BLEU-penalized" scores (indicated by BP-1 anad BP-4) which consider the BLEU score at a given test instance as 0 if the document retrieved for the given instance is incorrect. These help ensure that a model is not able to produce a high score by memorizing on a particular given domain.

4.4 Results

As can be seen in Table 1, VRAG outperforms RAG on each dataset for document retrieval. We also note that both RAG and VRAG significantly improve the performance of the initial DPR based retriever. Table 2 shows the performance of the models on the response generation task. Note that all results in Table 2 are obtained after further fine tuning the generator networks (after joint training is complete). The results show that the VRAG model outperforms the RAG model on language generation tasks (BLEU metrics) on all datasets except DoQA. We believe this difference is due to the nature of documents used in this dataset - other datasets have a lot more fact based questions to be answered using knowledge in documents, while more than 66% of the questions in the DoQA dataset are non-fact based and open-ended.

In addition, it is also possible that the RAGmodel is memorizing and overfitting on this dataset. For instance, see gains in penalized BLEU scores (BP-1 and BP-4) of RAG and VRAG over DPR in OR-QuAC and DSTC9 datasets in Table 2 - the relative gain of VRAG (over RAG) is significantly higher. This suggests that VRAG model is more likely to generate the response using the correct document and not by merely memorize on a given domain.

Alternative approximations for Variational Training: In Table 1 we see that the recall scores of pre-VM model are much worse than our VRAG model. This observation validates our hypothesis that, because the retrieved samples are not improved during training, the prior and posterior distributions in pre-VM end up focusing only on the few (potentially incorrect) initially-retrieved documents. This is especially problematic if the initial recall is low. As training progresses, this may worsen the distributions over the initial model (eg: see DSTC9 recall scores for DPR and pre-VM in Table 1) because the correct document isn't present in the retrieved sample, thus giving it an incorrect signal.

Effect of Top-5 Decoding: From Table 2 we find that in almost all cases, using top-5 decoding to generate responses performs better than using the single (best scored) document to generate responses. This indicates that models are to able to incorporate information from the correct document even if it is not returned as the top-ranked document.

Benefit of Document Posterior: In Section 3, we motivated the VRAG model by suggesting that using the posterior to sample documents while training the decoder could help train a better model. We find that recall of the document posterior in VRAG is nearly 14-70% higher than the document prior of RAG (depending on the dataset). Further, we find that the use of the responses by VRAG (posterior), to query the index during training, results in significantly better retrieval accuracy than RAG (prior) at every epoch during training. While this is perhaps intuitive and expected, our results on recall in Table 2 demonstrate the VRAG (prior) which is trained indirectly via the KL-divergence with VRAG (posterior) also outperforms the RAG (prior). We attribute this gain to the fact that the generator in VRAG learns to focus well on the generated documents which further reinforces the posterior network (and indirectly the prior network through KL term) to improve.

Study of Memorization: One of the issues in such an unsupervised learning is that generator may fail to use the retrieved knowledge. Instead the parameters might have been trained to internally model the external knowledge sources themselves, refered to as 'memorization'. In order to study memorization in the models, we compared the results on response generation of all models in the absence of the correct document - if the correct document is missing, the models should perform very poorly. A good performance even without obtaining correct document, would indicate lesser reliance on external knowledge and hence a higher tendency for 'memorization'. We, thus rebuild the document index without including any of the documents from the test set and then reevaluate the performance of our models. Table 3, shows the drop in BLEU scores for each of the models after removing the correct document on OR-QuAC dataset. As can be seen, the percentage drop is highest in the VRAG model indicating a higher usage of knowledge instance and thus, possibly lesser memorization.

5 Conclusion

In this paper we described an approach to run variational training on knowledge grounded dialog with a large corpus. Our experiments on three conversational QA datasets indicate that variational training is helpful as it produces better document samples while training. We find that our model, VRAG (having access to superior samples from posterior while training), not only generates better responses, it also learns a better retriever (prior distribution).

We believe that such sampling approximations could also be helpful in other tasks; for instance, it could also be interesting to apply them to other approaches such as Reinforcement Learning to the setting of a large corpus. This would require overcoming similar challenges in sampling as we did in this paper for variational training.