Learning Retrieval Augmentation for Personalized Dialogue Generation

Based on the PersonaChat dataset Zhang et al. (2018), Dinan et al. curate the CONVAI2 dataset, which contains a brief persona with four to five sentences for each interlocutor. This unique dataset has become a standard benchmark for the personalized dialogue generation task and built on this dataset, there are numerous studies, each of which approaches personalized dialogue generation from diverse perspectives. For example, Wolf et al. proposes a fine-tuning model based on the GPT2 model Radford et al. (2019). Song et al. integrates three BERT models Devlin et al. (2019) via reinforcement learning to generate responses. Liu et al. propose a transmitter and receiver model, which utilizes reinforcement learning with manually designed rewards for further refinement, for the personalized dialogue generation task.Cao et al. adopt model-agnostic data augmentation to use language models, such as GPT2 and BERT Devlin et al. (2019), to augment the training set with pseudo training data. Huang et al. devise an adaptive attention mechanism to integrate information from persona and context encoders seamlessly. In contrast to the aforementioned models, the proposed LAPDOG method introduces an end-to-end dense retriever framework to simultaneously augment the input of the generator from external data source and tune the retriever.

There are some works to incorporate retrievers into their respective models via different integration strategies to enhance text generation tasks. For instance, DocPrompting Zhou et al. (2022) curated a retrieval annotation dataset to train a retriever to retrieve and do input augmentation for code generation. Toolformer Schick et al. (2023) bootstraps retrieval-annotated data, which performs fine-tuning on language models for the retrieval-augmentation ability. FLARE Jiang et al. (2023) extends the Toolformer to large language models such as Ouyang et al. (2022) with special-designed prompts. RePlug Shi et al. (2023) further refines the retriever by distilling the knowledge from the language model’s probability. RAG Lewis et al. (2020) jointly trains the retriever and language model, which updates the retriever by the language model’s probability. Different from those models, the proposed LAPDOG model is designed specifically for personalized dialogue generation with a focus on optimizing desired objectives rather than the language model’s probability distribution. Since RePlug, Toolformer, and FLARE are based on large language models or their API calls, we do not include them in the comparison to LAPDOG. Compared with other models, we do not rely on retrieval annotations or bootstrapped retrieval annotations. The training objectives are directly computed from a comparison between the generated text and ground truth, rather than relying on training probabilities that are not always aligned with the desired metrics. Additionally, we introduce a candidate augmentation to avoid the limitations of a confined candidate set. This broadens the scope of potential dialogues and better captures the richness and diversity of an agent’s persona.

In a persona-based conversation session denoted by $C=\{P,U\}$, the persona $P=\{p_{1},\ldots,p_{e}\}$ consists of $e$ profile sentences providing background information about a machine interlocutor $m$ and the dialogue context $U=\{u_{h,1},u_{m,1},...,u_{h,n}\}$ encompasses the exchange of utterances between a human interlocutor $h$ and the machine interlocutor $m$. In the task of persona-based dialogue generation, the persona $P$ is used to characterize the machine interlocutor $m$, but it contains only four to five sentences, i.e., $4\leq e\leq 5$. The conversation always starts by the human interlocutor $h$. The primary objective of this task is to generate the response $r=u_{m,n}$ based on the given persona $P$ and dialogue context $U$.

The persona $P$ is short and hence cannot give a full characterization for the background information. To enrich the persona, we utilize a retrieval corpus $\mathcal{D}$ consisting of stories from a story dataset (e.g., ROCStory Mostafazadeh et al. (2016)). Note that there is no explicit annotation between $\mathcal{D}$ and $P$, necessitating an alternative approach to evaluate the usefulness of the retrieval content.

As shown in Figure 1, the architecture of the LAPDOG model consists of a generator, which adopts a transformer-based encoder-decoder structure, to generate dialogues and a dual-encoder retriever to efficiently obtain relevant information from an external story corpus.

It is straightforward to directly train the generator and retriever using the generator’s probability distribution in a way similar to the RAG method. However, this strategy does not work well since the retriever would trap into a fixed candidate set and the predicted probability distribution is not always aligned with the desired objectives in the personalized dialogue generation task. Hence, as depicted in Figure 1, the LAPDOG model adopts a two-stage training procedure. In the first stage, the training process starts with supervised training for the generator (refer to Section 3.3.1). In the second stage, the framework starts to tune the retriever and learn the retrieval augmentation jointly. To learn retrieval augmentation (refer to Section 3.3.2), the retriever’s loss is computed from the evaluation metrics between the output of the generator and the ground truth. During the process of learning retrieval augmentation, to prevent the retriever from stagnating around a limited set of candidates, we design the retrieval candidate augmentation (refer to Section 3.3.3), a method ensuring diversity in the retrieval process. Afterward, we enrich the input of the generator with retrieval-enhanced data and compute a generator loss based on the augmented input (refer to Section 3.3.4). Finally, we combine losses from both the generator and retriever to jointly train the two components (refer to Section 3.3.5). In the following sections, we introduce each part in detail.

During inference, stories from the ROCStory dataset are fetched in alignment with the provided persona and then integrated into the dialogue context using the Fusion-in-Decoder (FiD) technique Izacard and Grave (2021). Each combination of story, persona, and context is individually encoded. These encoded contexts are concatenated and processed via cross-attention in the decoder to produce the final response in an auto-regressive fashion. Additional experiments evaluating the effects of various query combinations, such as persona+context and context alone, are detailed in Appendix A.2 to highlight their impact on performance.

ConvAI2 Dinan et al. (2019) is a dialogue dataset collected from the crowd, featuring 8939/1000 multi-turn conversations that rely on 1155/100 persona descriptions for the train/dev splits. Each persona is succinctly depicted by approximately 5 profile sentences. Paired workers engaged in interactive conversations based on predefined personas.

Given the absence of a paired annotated retrieval corpus connected to ConvAI2, we employ ROCStory Mostafazadeh et al. (2016) as an auxiliary retrieval dataset. Our aim is for the narratives within this dataset to serve as supplemental content to the existing personas within the dialogue. We have undertaken pre-processing of the ROCStory to align the narrative style more closely with persona representation, including changes like transforming ‘he’ to ‘I’ and ‘does’ to ‘do’. The detailed pre-processing is listed in Appendix A.7. Statically, there are 98,161 stories within the corpus, and each story is composed of $5$ sentences.

We employ T5 series models Raffel et al. (2020) (small, base, XL) as the foundational model used for the generator. We initialize our generator with pre-trained weights from T5 and subsequently fine-tune it on the CONVAI2 dataset as $\mathcal{G}_{sup}$. The dense retriever is initialized with Contriever Izacard et al. (2021), a dual-encoder retriever that shares a similar encoder structure to BERT Devlin et al. (2019). We performed fine-tuning on both the retriever and generator using $\mathcal{G}_{sup}$, with a learning rate of $5\text{\times}{10}^{-4}$ and $\rho=0.5$. We further tune for learning retrieval augmentation based on the supervised foundation models in one epoch. We use persona profile as the query to retrieve the relevant stories.

LAPDOG aims to optimize for some generation metrics to enhance the dialogue quality. The evaluation comprises three metrics. The first metric is F1, which computes the harmonic mean of precision and recall on a word level between the generated text and the ground truth. The second metric is BLEU Papineni et al. (2002); Post (2018), an $n$-gram precision-based measure that quantifies the overlap between the generated text and the ground truth by penalizing for overly long or short outputs. The third metric is ROUGE-L Lin (2004), a variant of ROUGE that considers the longest common subsequence between the generated text and the ground truth, to effectively measure sentence-level structural similarity. With those metrics, we ensure a comprehensive assessment of the quality of generated dialogues. To enhance the aforementioned three metrics, we sum these three metrics together as the overall metric to train the LAPDOG model.

To compare the enhancement between different retrieval-augmentation approaches, we compared the results with the following baselines. First, we compare the LAPGOG with T5${}_{sup}^{S/B/XL}$ models, where $S/B/XL$ indicates the model size small, base, XL respectively. T5${}_{sup}^{S/B/XL}$ serves as the foundation models $\mathcal{G}_{sup}$. We also add a fixed retriever $\mathcal{R}_{fix}$ initialized from Contriever to validate the effectiveness of tuned and untuned retrievers. Meanwhile, we utilized the reinforcement learning tuning (T5${}^{S}_{Sup}$+$\mathcal{R}_{fix}$+RL) as one baseline, where the reward is set as the desired objective. Lastly, we introduce the RAG tuning that updates the retriever based on the generator’s training output probabilities instead of the desired metric, which is to validate the direct and indirect metric tuning for the objective.

The T5${}^{S}_{Sup}$ model forms our baseline. Augmenting it with a fixed retriever, T5${}^{S}_{Sup}$+$\mathcal{R}_{fix}$, shows a slight improvement in F1 and BLEU scores but a small decrease in the ROUGE-L score. This indicates a moderate enhancement in both F1 and BLEU, though the decrease in ROUGE-L suggests a trade-off in terms of capturing long-distance dependencies.

The results for reinforcement learning tuning, T5${}^{S}_{Sup}$+$\mathcal{R}_{fix}$+RL, exhibit a significant degradation across all metrics, indicating that this method might be not so effective for this task. This could be due to the challenge of setting an appropriate reward function for reinforcement learning.

The T5${}^{S}_{Sup}$+RAG model slightly outperforms the baseline model in terms of F1 but performs worse in terms of BLEU and ROUGE-L. This suggests that while the model seems to generate more correct words, there may be a compromise on the overall grammatical and semantic quality of the generated text. In contrast, the LAPDOG-enhanced model, T5${}^{S}_{Sup}$+LAPDOG, shows the highest improvements in all metrics among small models. This indicates that LAPDOG significantly enhances the ability to generate high-quality text and captures the desired metrics more effectively than other models.

For larger models, similar phenomena are observed. LAPDOG consistently delivers the best improvements over the base model, no matter whether it is T5${}^{B}_{Sup}$ or T5${}^{XL}_{Sup}$. This suggests that the efficacy of LAPDOG is not confined to smaller models and scales well with the model size.

To understand the individual contribution of each metric, we perform ablation experiments by successively removing one metric from the combined optimization process with results shown in Table 2. When we remove BLEU, the performance experiences a slight drop across all metrics with BLEU decreasing to 3.07, ROUGE-L to 14.35, and F1 to 14.29, suggesting that BLEU contributes to a more precise matching of response generation. Moreover, when we exclude F1, we see a more significant reduction in the performance, indicating the crucial role of F1 to ensure the overlap of words between the generated responses and the ground truth. Lastly, the removal of ROUGE-L also results in a decrease in the performance across all three metrics, showing its essential role in evaluating the coherence of generated dialogues. In summary, each metric contributes uniquely to the optimization process, and their combination in LAPDOG provides a more comprehensive guide for the generation of high-quality, personalized dialogues.

As shown in Figure 2, where $K$ increases from 2 to 6, we observe a consistent improvement in all metrics (i.e., BLEU, ROUGE-L, and F1). Generally, increasing the number of retrieval candidates improves the performance of the model, as evidenced by the improvements in the BLEU, ROUGE-L, and F1 scores. Interestingly, it is observed that the model performance does not monotonically increase with the number of candidates. The performance fluctuates as $K$ varies, implying that the number of retrieval candidates needs to be carefully selected. Too few candidates may not provide enough information for generating responses, while too many ones may introduce irrelevant information, which could potentially confuse the model.

The influence of candidate augmentation is explored in two aspects: quantitative performance and the diversity of retrieved stories. Table 3 provides a comparison between the performance of the LAPDOG model with and without candidate augmentation. The incorporation of candidate augmentation leads to superior performance across all three evaluation metrics. Specifically, LAPDOG without candidate augmentation attains slightly lower scores in all three metrics. This indicates that the inclusion of candidate augmentation enhances the overall performance of our model, confirming its crucial role in the proposed LAPDOG model.

To further investigate the impact of candidate augmentation to the retrieval diversity, we count the number of unique stories retrieved during testing. As shown in Figure 3, the LAPDOG method with candidate augmentation retrieves 1570 unique stories, whereas without it, the model only retrieves 738 unique stories. This result implies that candidate augmentation significantly contributes to the retrieval diversity. Other methods like RAG and fixed retriever manage to retrieve 25 and 1351 unique stories, respectively. This underlines the effectiveness of the proposed candidate augmentation approach in enhancing the diversity of the retrieval process, which in turn can help generate more personalized and contextually rich responses.

In contrast to the two-stage training process, we perform an ablation study by training LAPDOG from scratch, bypassing the first stage. As shown in Table 3, the results exhibit a significant decrease in all the metrics compared with the two-stage approach. Additionally, training directly from scratch requires more time to converge when compared with the two-stage training process. Overall, the two-stage training process is essential from both performance and efficiency standpoints.