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. Feb. 2026)

Abstract

Retrieval-Augmented Generation (RAG) improves Large Language Models (LLMs) by retrieving supporting documents into the prompt, but existing methods do not explicitly target queries that require fetching multiple documents with substantially different content. Such multi-aspect queries are challenging because relevant documents can be far apart in embedding space, making joint retrieval difficult. We introduce Multi-Head RAG (MRAG), which addresses this gap with a simple yet powerful idea: using Transformer multi-head attention activations rather than the standard decoder-layer embedding, as retrieval keys. It leverages the observation that different heads capture different semantic aspects. This yields multi-aspect embeddings for both documents and queries, improving retrieval accuracy on complex queries. We show MRAG's design advantages over 18 RAG baselines, up to 20% higher retrieval success ratios for real-world use cases, and improved downstream LLM generation. MRAG integrates seamlessly 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={2026},
  month={02},
  doi={10.48550/arXiv.2406.05085},
}