O ne G en: Efficient One-Pass Unified Generation and Retrieval for LLM s (2024)

Abstract: Despite the recent advancements in Large Language Models (LLMs), which have significantly enhanced the generative capabilities for various NLP tasks, LLMs still face limitations in directly handling retrieval tasks. However, many practical applications demand the seamless integration of both retrieval and generation. This paper introduces a novel and efficient One-pass Generation and retrieval framework (OneGen), designed to improve LLMs’ performance on tasks that require both generation and retrieval. The proposed framework bridges the traditionally separate training approaches for generation and retrieval by incorporating retrieval tokens generated autoregressively. This enables a single LLM to handle both tasks simultaneously in a unified forward pass. We conduct experiments on two distinct types of composite tasks, RAG and Entity Linking, to validate the pluggability, effectiveness, and efficiency of OneGen in training and inference. Furthermore, our results show that integrating generation and retrieval within the same context preserves the generative capabilities of LLMs while improving retrieval performance. To the best of our knowledge, OneGen is the first to enable LLMs to conduct vector retrieval during the generation.

## introduction in the era of large language models (llms), many natural language processing (nlp) tasks can be reduced to generation, allowing them to be addressed by a single llm @xcite @xcite . while llms excel in language generation, they still suffer from hallucinations (e.g., factual inaccuracies), stemming from their exclusive reliance on the parametric knowledge they contain @xcite @xcite . one promising approach is retrieval-augmented generation (rag) @xcite @xcite @xcite , which augments the input by retrieving relevant passages based on the query either before or during generation. other methods @xcite anchor llm generation to an external knowledge base through entity linking (el) during or after generation. these systems typically rely on a retriever at various stages of generation. however, due to the separate training paradigms for generation and retrieval, most prior work @xcite ) employs a separate model for text embedding. however, this pipeline approach has several drawbacks: i) deploying and maintaining two separate models introduces additional hardware overhead and increases maintenance costs. ii) the separation of models creates two distinct representational spaces, limiting interaction between the retriever and generator (e.g., llm) to text (i.e., query). as a result, whether the query is generated by the llm or input directly by the user, it requires an additional forward pass through the retriever, increasing inference computational costs. iii) in multi-turn dialogues, as illustrated in figure 1 (a), query rewriting is required for follow-up questions like "who is his wife?". this rewriting adds inference overhead and risks error propagation if inaccurate. iv) additionally, the pipeline approach is difficult to optimize end-to-end and requires large amounts of training data, while end-to-end optimization has been shown to yield significant benefits @xcite . our work introduces an efficient one-pass unified generation and retrieval (onegen) framework to enable an arbitrary llm to generate and retrieve in one single forward pass. inspired by the latest success in llm for text embedding @xcite , we expand the original vocabulary by adding special tokens (i.e. retrieval tokens) and allocate the retrieval task to retrieval tokens generated in an autoregressive manner. during training, retrieval tokens only participate in representation fine-tuning through contrastive learning @xcite , whereas other output tokens are trained using language model objectives. at inference time, we use retrieval tokens for efficient retrieving on demand. unlike previous pipeline approaches, which require at least two models for retrieval and generation (as shown in figure 1(b) ), onegen unifies both tasks into a single model, eliminating the need for a separate retriever. @xcite presents generative representational instruction tuning (grit), which aligns with this approach by training one llm to handle both generative and embedding tasks through different prompts and attention mechanisms, as depicted by the "switch" in figure 1 (c). however, grit still necessitates independent forward passes for generation and retrieval tasks, reducing efficiency for tasks that intertwine generation and retrieval. we evaluate the effectiveness of our method on two main tasks that require both generation and retrieval: rag (including single-hop qa which needs single-retrieval and multi-hop qa which needs multi-retrieval) and entity linking (el). empirical results show onegen outperforms the previous pipeline solutions as well as grit where applicable. specifically, onegen achieves +1.5pt improvement on average with four single-hop qa datasets on top of self-rag @xcite , +3.3pt f1 on average with two multi-hop qa datasets under three different 7b-based llms, and +3.2pt accuracy on average with 6 out-ofdomain entity linking datasets, with less training data. moreover, further analysis demonstrates one-gen can enhance retrieval capability when jointly trained, with no sacrifice in generation capability. in addition, we demonstrate superior inference speed and memory consumption of onegen compared with other llm alternatives, particularly as retrieval frequency increases. in summary, our work makes the following contributions: ## preliminaries and related works most text-based tasks can be reduced to generation, retrieval, or combination of the two. we first introduce several hybrid tasks and their common solutions in § 2.1. then, we introduce the three roles of tokens in llms in § 2.2. finally, we further explain the motivation of our method in § 2.3. ## onegen: one-pass generation and retrieval for llms we introduce a one-pass generation and retrieval framework (onegen) for fine-tuning llms on generation, retrieval, or hybrid tasks, as shown in figure 2 . our core idea is to integrate generation and retrieval to the same context by allocating the retrieval task to retrieval tokens generated in an autoregressive manner, thus enabling llm to perform both tasks in a single forward pass. ## conclusion and future work in this paper, we utilize mathematical notation to formally unify generative tasks, retrieval task, and their composites such as rag and el. for composite tasks, we integrate retrieval and generation within the same context. building upon this unified approach, we propose the onegen training framework, which harmonizes and expands both generative and representative instruction tuning. we conduct extensive experiments on two distinct types of composite tasks, rag and el, to validate the pluggability, effectiveness, and efficiency of onegen in training and inference. furthermore, our results confirm that integrating generation and retrieval within the same context does not negatively impact the generative capabilities of llms, while also providing significant enhancements in retrieval capabilities. future research directions include: 1) extending onegen to the multimodal domain for tasks such as multimodal rag and multimodal el. 2) enhancing onegen's training with diverse datasets to improve llms' capability for complex retrieval and generation tasks.