M - RAG : Reinforcing Large Language Model Performance through Retrieval-Augmented Generation with Multiple Partitions (2024)

Abstract: Retrieval-Augmented Generation (RAG) enhances Large Language Models (LLMs) by retrieving relevant memories from an external database. However, existing RAG methods typically organize all memories in a whole database, potentially limiting focus on crucial memories and introducing noise. In this paper, we introduce a multiple partition paradigm for RAG (called M-RAG), where each database partition serves as a basic unit for RAG execution. Based on this paradigm, we propose a novel framework that leverages LLMs with Multi-Agent Reinforcement Learning to optimize different language generation tasks explicitly. Through comprehensive experiments conducted on seven datasets, spanning three language generation tasks and involving three distinct language model architectures, we confirm that M-RAG consistently outperforms various baseline methods, achieving improvements of 11%, 8%, and 12% for text summarization, machine translation, and dialogue generation, respectively.

## introduction introduced by @xcite , retrieval-augmented generation (rag) represents a paradigm within the domain of large language models (llms) to augment generative tasks. more specifically, rag incorporates an initial retrieval step where llms query an external database to acquire relevant information before progressing to answer questions or generate text. this process not only guides the subsequent generation step but also guarantees that the responses are firmly anchored in the retrieved information (referred to as memories). consequently, it enhances llm performance, and has attracted growing research interests @xcite in recent years. while the majority of existing studies @xcite @xcite adopt a retrieval approach that considers a database as a whole, which tends to yield a coarse-grained retrieval. the collective organization of all memories may hinder the focus on crucial memories and introduce noise, particularly due to the inherent challenges of approximate k-nearest neighbor (aknn) search when applied to large datasets. in this context, we investigate a retrieval approach that aims to search within a partition of the database, corresponding retrieval at a fine-grained level, which is designed to enhance the generation process by targeting specific memories. moreover, in quite a few vector database systems, database partitions are regarded as fundamental units for analysis. this facilitates the construction and maintenance of index structures @xcite , ensures the protection of user privacy data (stored in specific partitions with access rights) @xcite , and supports distributed architectures @xcite . therefore, in this work, we propose to take a partition as a basic entity in the execution of rag, which is less explored in current methods. we discuss our proposal with a motivating experiment illustrated in figure 1 . we investigate various strategies for partitioning a database (elaborated in section 3.1), and perform rag with varying the number of partitions for three generation tasks: summarization, translation, and dialogue generation, where we explore all partitions for the retrieval, and the best result (assessed based on a development set) across different partitions is reported. we observe that the optimal performance is typically not achieved through retrieval based on the entire database (#@xmath0. this observation inspires us to investigate a novel rag setting with multiple partitions. to achieve this, the task should address three significant challenges, summarized below. (1) determining a strategy for partitioning a database and the number of partitions. our contributions can be summarized as follows: (1) we propose a multiple partition paradigm for rag, aiming to facilitate fine-grained retrieval and concentrate on pivotal memories to enhance overall performance. in addition, the utilization of multiple partitions benefits other aspects of rag, including facilitating the construction and maintenance of indices, protecting user privacy data within specific partitions, and supporting distributed parallel processing across different partitions. (2) we introduce m-rag, a new solution based on multiagent reinforcement learning that tackles the three challenges in executing rag across multiple partitions. we show that the training objective of m-rag is well aligned with that of text generation tasks. ## related work retrieval-augmented generation. we review the literature of retrieval-augmented generation (rag) in terms of (1) naive rag, (2) advanced rag, and (3) modular rag. for (1), naive rag follows a standard process including indexing, retrieval, and generation @xcite . however, its quality faces significant challenges such as low precision, hallucination, and redundancy during the process. for (2), advanced rag is further developed to overcome the shortcomings of naive rag. specifically, during the indexing stage, the objective is to enhance the quality of the indexed content by optimizing data embedding @xcite . during the retrieval stage, the focus is on identifying the appropriate context by calculating the similarity between the query and chunks, where the techniques involve fine-tuning embedding models @xcite , or learning dynamic embeddings for different context @xcite . during the generation stage, it merges the retrieved context with the query as an input into large language models (llms), where it addresses challenges posed by context window limits with re-ranking the most relevant content @xcite , or compressing prompts @xcite . in addition, self-rag @xcite is proposed to identify whether retrieval is necessary, or the retrieved context is relevant, which helps language models to produce meaningful generation @xcite . for (3), modular rag diverges from the traditional naive rag structure by incorporating external modules to further enhance the performance, including search module @xcite , memory module @xcite , tuning module @xcite , and task adapter @xcite . specifically, selfmem @xcite incorporates a retrieval-enhanced generator to iteratively create a memory pool, it then trains a selector to choose one of the memories from the pool to generate responses. the work @xcite provides a comprehensive survey of rag for llms. our work differs from existing rag studies in two aspects. first, we introduce a multiple partition setting, where each partition serves as a fundamental entity for retrieval, rather than retrieving from the entire database. second, we introduce an m-rag framework built upon multiagent reinforcement learning, which tackles three distinct challenges posed by this novel setting. reinforcement learning for llms. recently, reinforcement learning has seen broad applications across a variety of language-related tasks for large language models (llms). this includes tasks such as text summarization @xcite , machine translation @xcite , dialogue systems @xcite , semantic parsing @xcite , and review generation @xcite . for example, webgpt @xcite incorporates a reinforcement learning framework to autonomously train the gpt-3 model using a search engine during the text generation process. further, instructgpt @xcite collects a dataset containing desired model outputs provided by human labelers. subsequently, it employs reinforcement learning from human feedback (rlhf) to fine-tune gpt-3 @xcite . in addition, r3 @xcite introduces a rewrite-retrieve-read process, where the llm performance serves as a reinforcement learning incentive for a rewriting module. this approach empowers the rewriter to enhance retrieval queries, consequently improving the reader's performance in downstream tasks. mmqs @xcite introduces a new multimodal question suggestion task with a multi-agent version of rlhf. in this work, we propose a novel multi-agent reinforcement learning framework utilizing two agents to collaboratively optimize text generation tasks. to our best knowledge, this is the first of its kind. ## methodology a task involving m-rag can be formulated below. given a database @xmath1, y i )} |d| @xmath2.g., summarization), where each pair (x, y) represents a document and its corresponding summary stored in d. the m-rag initiates the process by partitioning d into multiple partitions. this can be achieved through meth- ## conclusion and limitations in this paper, we propose a multiple partition paradigm for rag, which aims to refine retrieval processes and emphasize pivotal memories to improve overall performance. additionally, we introduce m-rag, a novel framework with multi-agent reinforcement learning, which addresses key challenges inherent in executing rag across multiple partitions. the training objective of m-rag is well aligned with that of text generation tasks, showcasing its potential to enhance system performance explicitly. through extensive experiments conducted on seven datasets for three language generation tasks, we validate the effectiveness of m-rag. for limitations, we conduct experiments with quantized versions of language models due to computational constraints. however, the observed effectiveness gains are expected to remain consistent across different model sizes and should not significantly impact the overall trends of various rag methods. further, although the parameters of the llms remain fixed and only the parameters of agent-s and agent-r are trained, the training efficiency is limited, as indicated by the training time complexity discussed in section 3.4. this is due to the necessity of querying the llms during the training process. in future work, we intend to explore solutions to overcome these limitations. overall, we observe that m-rag consistently outperforms selfmem across diverse translation datasets, as evidenced by various evaluation metrics.