R ega VAE : A Retrieval-Augmented G aussian Mixture Variational Auto-Encoder for Language Modeling (2023)

Abstract: Retrieval-augmented language models show promise in addressing issues like outdated information and hallucinations in language models (LMs). However, current research faces two main problems: 1) determining what information to retrieve, and 2) effectively combining retrieved information during generation. We argue that valuable retrieved information should not only be related to the current source text but also consider the future target text, given the nature of LMs that model future tokens. Moreover, we propose that aggregation using latent variables derived from a compact latent space is more efficient than utilizing explicit raw text, which is limited by context length and susceptible to noise. Therefore, we introduce RegaVAE, a retrieval-augmented language model built upon the variational auto-encoder (VAE). It encodes the text corpus into a latent space, capturing current and future information from both source and target text. Additionally, we leverage the VAE to initialize the latent space and adopt the probabilistic form of the retrieval generation paradigm by expanding the Gaussian prior distribution into a Gaussian mixture distribution. Theoretical analysis provides an optimizable upper bound for RegaVAE. Experimental results on various datasets demonstrate significant improvements in text generation quality and hallucination removal.

## introduction language models (lms) have achieved state-ofthe-art performance on many nlp tasks @xcite , which reveals that they store a large amount of world knowledge as implicit parameters. while this development is exciting, lms still suffer from some problems @xcite : 1) performance and model parameter size follow a power law relationship (kaplan model future info. in aggregation query key value knn-lm explicit rag explicit realm explicit spalm explicit fid implicit emdr 2 implicit epr explicit re2g implicit retro implicit regavae implicit recently, the advent of retrieval-augmented text generation has emerged as a novel paradigm aimed at addressing these pertinent issues @xcite @xcite . compared to generative-only models, this paradigm not only explicitly exploits similar texts to generate more fluent sentences but also leverages expertise to generate difficult responses. nonetheless, we contend that there are two primary challenges associated with current retrieval-augmented language models. firstly, not only current semantic information, but also future semantic information need to be considered during retrieval. previous studies @xcite @xcite either directly use the entire text as key and value parts at the same time, and then use cosine similarity @xcite , tf-idf and other indicators to search, which leads to the value part is only similar to the source text (query), and does not necessarily serve the best for generator. another way is to divide a piece of text into two parts, where the first part and the second part are regarded as current information and future information, such as retro @xcite . however, retro adds future information to value part, but ignores the future information in query and key, which leads to the fact that candidate documents with high similarity do not necessarily contain future information that can help the generator. secondly, explicitly aggregating retrieved documents and source texts is limited by the length of the model input and introduces too much noise. implicit aggregation is inefficient in irregular embedding spaces, and retrieval vectors are not generalizable. to address the above challenges, we design re-gavae, a retrieval-augmented language model based on gaussian mixture variational auto-encoder. unlike previous methods that directly encode unlabeled corpora @xcite or only adding future information to the value part @xcite , as shown in tab. 1, our model considers future information through a latent space, given an x, we decode it into a y using a conditional vae, which ensures that the latent variables contain information from both source and target data. in addition, in order to implicitly aggregate the retrieved documents and source texts, we also use the probabilistic form of the retrieval generation paradigm to theoretically extend the prior gaussian distribution to a gaussian mixture distribution. this allows the latent space to satisfy continuity and uniformity, and the latent vector after aggregating retrieved documents and source text has better representation ability. tab. 1 summarizes the differences between regavae and existing representative methods. overall, our contributions are as follows: ## related work we classify related studies into two categories, explicit aggregation and implicit aggregation, according to the way the retrieved documents and source text are aggregated. explicit aggregation refers to concatenating retrieved documents directly into source text to construct augmented input. implicit aggregation refers to adding retrieved documents to the generator in the form of vectors or distributions. ## methodology most text generation tasks can be formulated as a mapping from a source text x to a target text y : @xmath0, while retrieval-augmented text generation can be further formulated as: @xmath1, r), where r is the relevant document retrieved based on x. specifically, this approach generally encompasses the utilization of a retriever denoted as r and a generator denoted as g. the retriever r obtains r from the retrieval source s by the retrieval metric d and x. then r and x are fed into g to obtain y through a predefined integration method i. commonly used retrieval indexes d include cosine similarity, tf-idf, etc. this paradigm can also be expressed in probabilistic form: next, the framework of regavae is introduced, which consists of three steps. firstly, in order to construct a compact space, we introduce the vae structure. since transformers based on vae all suffer from the posterior collapse @xcite , we follow a previous study @xcite which combines low-rank tensor products for latent variables and decoders (see § 3.1 and step 1 in fig. 1 ). secondly, to introduce retrieval information into the latent space, we first introduce how the retrieval library is constructed (see step 2 in fig. 1 ), and then replace the prior gaussian distribution in the original vae with a gaussian mixture distribution to derive regavae (see step 3 in fig. 1 ). this allows for deep aggregation of retrieved and input documents and simultaneously incorporates future information into query, key and value parts, which helps to generate more fluent sentences (see § 3.2) . finally, to train regavae, we derive an optimizable upper bound on the loss function for unclosed solutions (see § 3.3). fig. 1 shows the whole framework diagram. ## experiment this section provides the experimental datasets, experimental settings, and experimental results. ## analysis to further analyze regavae, we explore the impact of the number of retrieved neighbors, different model structures on the model performance. we also give a case to verify the model performance. ## conclusion in this paper, we summarize two major challenges of existing retrieval-augmented language model methods, and propose regavae to address them. we find that regavae outperforms traditional retrieval generative models in terms of both generative quality and reduce hallucinations. in addition, ablation experiments and three analysis experiments verify the correctness of the model motivation. in future work, we will consider migrating regavae to large language models.