Retrieval Augmented Generation (RAG) enhances the abilities of Large Language Models (LLMs) by enabling the retrieval of documents into the LLM context to provide more accurate and relevant responses. Existing RAG solutions do not focus on queries that may require fetching multiple documents with substantially different contents. Such queries occur frequently, but are challenging because the embeddings of these documents may be distant in the embedding space, making it hard to retrieve them all. This paper introduces Multi-Head RAG (MRAG), a novel scheme designed to address this gap with a simple yet powerful idea: leveraging activations of Transformer's multi-head attention layer, instead of the decoder layer, as keys for fetching multi-aspect documents. The driving motivation is that different attention heads can learn to capture different data aspects. Harnessing the corresponding activations results in embeddings that represent various facets of data items and queries, improving the retrieval accuracy for complex queries. We provide an evaluation methodology and metrics, synthetic datasets, and real-world use cases to demonstrate MRAG's effectiveness, showing improvements of up to 20% in relevance over standard RAG baselines. MRAG can be seamlessly integrated with existing RAG frameworks and benchmarking tools like RAGAS as well as different classes of data stores.

翻译：检索增强生成（RAG）通过将文档检索引入大语言模型（LLM）上下文，增强其生成更准确和相关性响应的能力。现有RAG方案未聚焦于需要获取内容差异显著的多个文档的查询。此类查询频繁出现，但由于这些文档的嵌入向量在嵌入空间中可能相距较远，导致难以全部检索。本文提出多头检索增强生成（MRAG）这一新方案，通过一个简单而强大的思想解决该问题：利用Transformer多头注意力层的激活值（而非解码器层）作为键值来获取多方面文档。核心动机在于不同注意力头能学习捕获数据的不同方面。利用对应激活值生成的嵌入向量能表征数据项和查询的多个侧面，从而提升复杂查询的检索准确率。我们设计了评估方法体系与指标、合成数据集及真实用例，证明MRAG的有效性——与标准RAG基线相比，相关性提升最高达20%。MRAG可无缝集成至现有RAG框架（如RAGAS评估工具）及各类数据存储系统。