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

TL;DR

This paper introduces M-RAG, a retrieval-augmented generation framework using multiple database partitions and multi-agent reinforcement learning, significantly improving performance across various language generation tasks.

Contribution

The paper proposes a novel multiple partition paradigm for RAG and a reinforcement learning framework to optimize language generation, outperforming existing methods.

Findings

M-RAG improves performance by up to 12% across tasks.

Consistent outperformance over baseline methods in experiments.

Effective across multiple datasets and language models.

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.