Copyright Notice:

The documents distributed by this server have been provided by the contributing authors as a means to ensure timely dissemination of scholarly and technical work on a noncommercial basis. Copyright and all rights therein are maintained by the authors or by other copyright holders, notwithstanding that they have offered their works here electronically. It is understood that all persons copying this information will adhere to the terms and constraints invoked by each author's copyright. These works may not be reposted without the explicit permission of the copyright holder.

Publications of SPCL

M. Besta, A. Kubicek, R. Gerstenberger, M. Chrapek, R. Niggli, P. Okanovic, Y. Zhu, P. Iff, M. Podstawski, L. Weitzendorf, M. Chi, J. Gajda, P. Nyczyk, J. Müller, H. Niewiadomski, T. Hoefler:

 Multi-Head RAG: Solving Multi-Aspect Problems with LLMs

(arXiv:2406.05085. Jun. 2025)

Abstract

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 observation is that different attention heads 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, multi-aspect datasets, and real-world use cases to demonstrate MRAG's effectiveness. We show MRAG's design advantages over 18 RAG baselines, empirical improvements of up to 20% in retrieval success ratios, and benefits for downstream LLM generation. MRAG can be seamlessly integrated with existing RAG frameworks and benchmarks.

Documents

download article:
access preprint on arxiv:
 

BibTeX

@article{besta2024multi,
  author={Maciej Besta and Ales Kubicek and Robert Gerstenberger and Marcin Chrapek and Roman Niggli and Patrik Okanovic and Yi Zhu and Patrick Iff and Michal Podstawski and Lucas Weitzendorf and Mingyuan Chi and Joanna Gajda and Piotr Nyczyk and Jürgen Müller and Hubert Niewiadomski and Torsten Hoefler},
  title={{Multi-Head RAG: Solving Multi-Aspect Problems with LLMs}},
  journal={arXiv:2406.05085},
  year={2025},
  month={06},
  doi={10.48550/arXiv.2406.05085},
}