PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Maciej Besta, Ales Kubicek, Robert Gerstenberger, Marcin Chrapek, Roman Niggli, Patrik Okanovic, Yi Zhu, Patrick Iff, Michal Podstawski, Lucas Weitzendorf, Mingyuan Chi, Joanna Gajda, Piotr Nyczyk, Jürgen Müller, Hubert Niewiadomski, Torsten Hoefler

arXiv:2406.05085v6cs.CLcs.AIcs.IR

TL;DR

Existing RAG methods do not explicitly target queries requiring retrieval of multiple semantically distinct documents, which may be far apart in embedding space. MRAG uses decoder-model attention-head activations for multi-aspect embeddings and reports up to 20% better retrieved-document relevance than modern RAG baselines, alongside improved downstream generation.

  • Problem

    Existing RAG solutions do not explicitly target multi-aspect queries, where relevant documents can be far apart in embedding space and difficult to retrieve together.

    Source: §1 · p. 1, §7 · p. 9

  • Method

    MRAG uses decoder-model multi-head attention activations instead of traditional feed-forward-layer embeddings to represent distinct aspects without additional training.

    Source: §1 · p. 2, §7 · p. 9

  • Results

    MRAG reports up to 20% better retrieved-document relevance than modern RAG baselines and improved downstream LLM generation.

    Source: §7 · p. 10

  • Takeaways & Limitations

    MRAG avoids additional LLM queries, model instances, storage, training, or multiple inference passes over the embedding model.

    Source: §7 · p. 10

    MRAG assumes pretrained attention-head embeddings already encode multi-aspectuality suitable for downstream retrieval, without additional fine-tuning.

    Source: §2.4 · p. 4

Abstract

from arXiv · show

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.

1 Introduction

Existing RAG does not explicitly target queries requiring semantically distinct documents that may be far apart in embedding space. MRAG uses attention-head activations to represent multiple aspects.

1 Introduction

  • Multi-aspect queries require integrating semantically distinct documents, which can lie far apart in embedding space and complicate joint retrieval.The introduction notes such queries in industrial settings, including accident analysis spanning worker psychology, equipment records, weather, or microclimate.Source: §1 · pp. 1–2
  • MRAG uses decoder-block multi-head attention activations as embedding sources, representing distinct input aspects without additional space or compute costs.The design requires neither additional fine-tuning nor architectural modifications to the base model.Source: §1 · p. 2
  • MRAG reports up to 20% higher retrieved-document relevance than modern RAG baselines across comparisons with 18 traditional and modern designs.Its evaluation also introduces multi-aspect datasets and a benchmarking methodology.Source: §1 · p. 2
  • MRAG maintains comparable performance on single-aspect queries, benefits downstream LLM generation, and integrates with existing RAG frameworks.The paper also says its benchmarking methodology can extend benchmarks such as RAGAs.Source: §1 · p. 2

2 MRAG: Design & Implementation

MRAG extends a standard RAG pipeline by storing head-specific embeddings and using them for multi-aspect retrieval, ranking, and generation.

2.1 Constructing Multi-Aspect Embeddings

  • MRAG stores h head-specific embeddings per item, preserving aspect distinctions while linking each vector to the original chunk; head importance combines embedding norm and semantic spread.Each embedding space receives a one-time score based on average L2 norm and average cosine distance.Source: §2.1 · pp. 2–3

2.3 Query Execution

  • Query vectors are computed in one pass; per-head top-c searches are merged by importance-weighted voting, then top-k chunks enter separate prompt sections.The voting score discounts lower-ranked candidates exponentially using their position in each head-space list.Source: §2.3 · pp. 3–4

2.4 Multi-Aspectuality Without Additional Training

  • MRAG assumes pretrained attention heads already encode useful aspect distinctions; prior literature and the authors’ attention-pattern analysis motivate avoiding additional fine-tuning.The analysis reports shifting token focus across heads.Source: §2.4 · p. 4

2.5 Parallelization and Lightweight System Design

  • All head vectors come from one forward pass, and parallel ANN searches keep compute and storage at standard RAG’s O(nd).The design uses off-the-shelf indexes and fixed-size vectors for parallel subspace search.Source: §2.5 · p. 4

3 Compute & Storage Complexity Analysis

The analysis compares MRAG’s runtime and storage complexity with 18 RAG baselines.

3 Compute & Storage Complexity Analysis

  • MRAG matches standard RAG latency and embedding dimensionality while requiring one preprocessing pass and simple head statistics; Poly-encoder and ColBERT incur higher costs.All h attention-head embeddings are extracted in parallel from a single forward pass.Source: §3 · p. 4

4 Benchmarking Multi-Aspectuality

The section defines datasets and retrieval metrics for evaluating whether RAG retrieves multiple aspects of a query.

4 Benchmarking Multi-Aspectuality

  • The benchmark includes Wikipedia, legal-document, and industry accident-report datasets, with queries combining 1–25 distinct categories.A query with 10 aspects, for example, asks about 10 distinct documents from 10 categories.Source: §4 · p. 4
  • Retrieval Success Ratio measures exact-document matches; Category Retrieval Success Ratio accepts same-category matches, and Weighted Success Ratio tunes the trade-off between them.The weighted metric combines exact and category-level success ratios using w.Source: §4 · p. 4
  • Category-level matches are included because same-category documents may still enhance generation when exact matches are missing.The motivation draws on prior work and the authors’ experience in legal and industrial accident analysis.Source: §4 · pp. 4–5
  • Table 1 compares time and storage complexity across RAG schemes; its caption states that only MRAG matches vanilla RAG asymptotically.Source: Table 1 · p. 5

5 Evaluation

Across multi-aspect benchmarks, MRAG improves retrieval over Standard RAG, including exact and category matches, while retaining single-aspect performance. It also benefits downstream generation without latency or storage overhead.

5 Evaluation

  • Evaluation compares Standard, Split, and Fusion RAG using SFR and e5, weighted retrieval success, and 25-query samples.The weighted ratio gives exact document matches twice the weight of category matches.Source: §5 · pp. 5–6

5.1 Superior Performance for Multi-Aspect Queries

  • >10% increase in exact-match retrieval success and >25% in category-match success; weighted performance improvement is 10–20% across aspect counts and SFR/e5.MRAG also outperforms Split RAG; the paper reports advantages on legal-document synthesis and accident-cause analysis.Source: Figure 3 · p. 6, §5.1 · pp. 6–7, Figure 4 · p. 7

5.2 High Performance Also for Single-Aspect Queries

  • MRAG performs on par with Standard RAG on single-aspect queries across the datasets and embedding models tested.Table 2 reports exact-document retrieval success for 25 single-aspect queries per dataset.Source: Table 2 · p. 7, §5.2 · p. 7

5.3 Seamless Enhancement of Existing RAG Schemes

  • Fusion MRAG consistently outperforms Fusion RAG, although both Fusion strategies add variance, compute, latency, and token costs.Figure 6 compares MRAG with Split RAG and Fusion MRAG with Fusion RAG.Source: §5.3 · p. 7, Figure 6 · p. 8

5.6 Confirmation of Specialization of Heads

  • Some head–category pairs retrieve more than 2× higher than average; for the Wikipedia multi-aspect dataset, mathematical-transforms retrievals dominate head 7, which Figure 7 identifies as the highest bar.Source: §5.6 · p. 8, Figure 7 · p. 8

5.4 Advantages for Downstream LLM Generation

  • MRAG retrieval improves downstream LLM generations, which contain more diverse and complementary facts for complex queries.The paper attributes richer responses to the multi-aspect embeddings increasing the likelihood of including such facts.Source: §5.4 · p. 8

5.5 No Latency and Storage Overheads

  • MRAG adds no query-embedding latency or storage compared with standard RAG; head-level extraction and retrieval are parallelizable.The total embedding dimensionality remains unchanged.Source: §5.5 · p. 8

5.7 Further Analyses & Ablation Studies

  • Additional analyses of decoder blocks, preprocessing scalability, and reranking strategies broadly confirm MRAG’s outperformance of other baselines.These analyses are reported in Appendix I.Source: §5.7 · p. 8

6 Related Work & Advantages of MRAG

MRAG uses decoder-model attention-head embeddings to represent distinct semantic aspects, unlike prior approaches that either retrieve multiple documents without multi-aspect embeddings or use token-level vectors. Its design also aims to keep preprocessing and inference practical while remaining compatible with existing retrieval components.

Practical advantages

  • Head scoring is parallelizable, preprocessing computes global embeddings in one document pass, and inference adds no cost over standard RAG.The MHA-derived vectors match standard last-layer embedding dimensionality, keeping storage overhead minimal.Source: §6 · p. 8

Comparison scope

  • The comparison of RAG schemes covers training needs, MHA compatibility, modality extensibility, inference overhead, preprocessing, storage, and multi-aspectuality.The table caption describes these comparison dimensions, not the values or winners.Source: Table 3 · p. 9

Compatibility

  • MRAG can serve as the dense branch in hybrid retrieval and can be combined with existing cross-encoder rerankers without changing the retriever or index.It is orthogonal to sparse retrieval and reranker choice.Source: §6 · p. 9

Related work and novelty

  • MRAG is presented as the first method to use MHA embeddings from pretrained decoder-based embedding LLMs for RAG, without architectural changes, training, or custom modules.Other multi-vector approaches use token-level embeddings to cover document content, rather than representing distinct semantic aspects.Source: §6 · pp. 8–9

7 Conclusion

Multi-aspect queries are difficult for standard RAG because relevant documents can be far apart in embedding space. MRAG uses separate attention-head activations to represent different aspects, and the conclusion reports improved retrieval and downstream generation.

7 Conclusion

  • Up to 20% better performance than modern RAG baselines, with improved downstream LLM generation, is reported for MRAG.The evaluation covers multi-aspect queries and reports richer, more factual LLM outputs.Source: §7 · p. 10
  • MRAG requires no extra LLM queries, model instances, storage, training, or embedding-model inference passes, unlike other RAG schemes.The conclusion presents these design properties alongside retrieval accuracy and scalability as advantages.Source: §7 · p. 10

A Review of Multi-Aspectuality in Industry

  • Safety investigations, integrated medical diagnostics, and enterprise decisions combine semantically distinct information sources because no single stream suffices for complex tasks.The review notes that standard RAG embeddings for divergent subdomains can be far apart.Source: A Review of Multi-Aspectuality in… · p. 19

B Mathematical Formulation: Additional Details

  • The formulation describes attention-head outputs and their combination into the multi-head attention block, with the output projection combining the heads.It omits layer-normalization details.Source: B Mathematical Formulation: Additional… · p. 19

B.1 Standard RAG Embedding

  • Standard RAG represents a chunk with one embedding from the final token’s decoder output after the final feed-forward layer.Source: §B.1 · p. 19

B.2 Multi-Head RAG Embedding

  • MRAG retains each attention head’s final-token output as a separate embedding, yielding h vectors per chunk that capture diverse semantic aspects without extra inference memory or compute.Source: §B.2 · p. 19

C System Design & Implementation: Additional Details

The supplied system and implementation details describe head scoring, chunk reranking, and data-store integration, alongside prompts for multi-article reasoning and synthetic query generation. The accompanying discussion motivates extracting attention-head representations without fine-tuning.

C.1 Ranking Strategy Details

  • MRAG scores each head by multiplying its average embedding norm by its average cosine distance to a uniformly sampled subset of embeddings.This procedure is presented as the head-importance scoring strategy.Source: §C.1 · p. 20

C.2 Ranking Strategy Details

  • The reranker collects each chunk’s best-matching chunks, sorts the resulting chunk-weight tuples, and returns the top k.Source: §C.2 · p. 20

C.3 Integration with Data Stores

  • MRAG supports exact or approximate nearest-neighbor search and different data stores; these choices are orthogonal to its retrieval method.Source: §C.3 · p. 20

Prompt Template for Reader

  • The reader prompt asks for a coherent answer that compares, contrasts, or synthesizes relevant claims across articles and supports it with specific references.Source: §D.1 · p. 21

D.2 Prompt Template for the Synthetic Dataset Generation

  • The synthetic-data prompt requests a story in which every attached article is substantively relevant, each title appears, and each article is referenced in at least two sentences.Source: §D.2 · p. 21

E Multi-Aspectuality with Attention Heads without Additional Training

E.1 Literature Survey

  • Prior studies report attention heads specializing in linguistic or token patterns, while pruning results indicate that many heads contribute little and specialization may concentrate in fewer effective heads.The cited survey includes findings from both GPT-style and BERT-style models.Source: §E.1 · p. 22

E.2 Analysis of Multi-Head Patterns

In LLaMA-2 7B and SFR-Embedding-Mistral, attention patterns differ across heads. The analysis describes recurring token-attention patterns and a line-break attention spike in the embedding model.

E.2 Analysis of Multi-Head Patterns

  • Attention patterns vary substantially across heads in both examined models, with diagonal patterns attending to tokens a constant number of steps earlier and horizontal or vertical patterns focusing on specific tokens.The authors postulate that diagonal patterns may help identify language rhythm and semantic connections; they suggest horizontal and vertical patterns may reflect focus on semantic aspects.Source: §E.2 · p. 22
  • When processing the final input token, different heads attend differently to many passages, although some tokens receive substantial attention from most heads.Source: §E.2 · p. 22
  • In SFR-Embedding-Mistral, all heads’ attention spikes positively or negatively on the first line break; the authors conjecture this reflects the model’s retrieval-instruction embedding usage.Source: §E.2 · p. 22

F Complexity Analysis: Additional Details

The appendix details how RAG complexity is accounted for and specifies the multi-aspect benchmark, including its datasets and query construction.

F Complexity Analysis: Additional Details

  • Complexity accounting models system-specific overheads: RQ-RAG adds sequential query work, ThinkNote triples work and doubles depth, and ColBERT scales retrieval and storage with query and document size.Meta-chunking also increases preprocessing, storage, and retrieval costs through chunking and additional vectors.Source: F Complexity Analysis: Additional… · pp. 22–24

G Full Specification of Benchmarking Multi-Aspectuality

  • The benchmark retrieves the top k documents for each query and checks them against the ground-truth documents and their non-overlapping categories.The definition of k includes all fetched documents, including mismatches.Source: G Full Specification of Benchmarking… · p. 24

G.1 Multi-Aspect Datasets

  • The datasets combine synthetic Wikipedia collections with real-world-inspired NDA and chemical-plant accident reports, totaling 16,500 documents.The real-world-inspired datasets were developed with tech leaders from three corporations that use RAG in internal report-generation and analytics frameworks.Source: §G.1 · pp. 24–25

G.2 Multi-Aspect Query Generation

  • The query set contains 160 queries at each of ten aspect counts, from 1 to 25, with each query combining documents sampled from distinct categories.An LLM generates a story that combines the sampled documents.Source: §G.2 · p. 25

G.3 Metrics

The metrics measure exact-document retrieval, retrieval from the correct categories, and a weighted combination of those outcomes.

G.3 Metrics

  • Retrieval Success Ratio measures the fraction of ideal relevant documents that are successfully retrieved.The ideal set is compared with the set actually retrieved for the query.Source: §G.3 · p. 25
  • Category Retrieval Success Ratio counts retrieved documents belonging to the categories of the ideal documents, even when they are not exact matches.It uses the same ratio form as Retrieval Success Ratio.Source: §G.3 · p. 25
  • Weighted Retrieval Success Ratio combines exact-document and category-level retrieval success.The supplied passage names the combined measure but does not provide a complete weighting formula.Source: §G.3 · p. 26

H Evaluation Setup: Additional Details

Additional experiments compare decoder-block and voting-strategy choices, and assess preprocessing overhead. The appendix also reports compute resources and points to a dataset overview.

H.1 Compute Resources

  • Experiments used nodes with four NVIDIA GH200 GPUs and 800 GB of total memory; one GPU with at least 40 GB should generally suffice.The main experiments took at most three hours of GPU time and up to $15 in OpenAI API costs.Source: §H.1 · p. 26

H.2 Dataset Details

  • Table 4 overviews the structure and document counts of the datasets used.Source: §H.2 · p. 26

I.1 Harnessing Different Decoder Blocks

  • The final multi-head attention block performs best against Standard RAG among the three decoder-block choices tested.The alternatives were the first attention block and block 16.Source: §I.1 · p. 26

I.2 Analyzing Different Voting Strategies

  • Both alternative MRAG voting strategies perform worse on average than the selected strategy.The alternatives use exponential significance lowering alone or chunk-to-query distance-based weights.Source: §I.2 · p. 26
  • The two alternative Split RAG voting strategies perform on par with each other but worse than MRAG.One uses exponential significance lowering; the other weights chunks by significance and position.Source: §I.2 · p. 26

I.3 Analyzing Preprocessing Overhead

  • On ArguAna, importance scoring adds under 5% to encoding time despite 310 seconds of total encoding; sampling makes the overhead lower than full pairwise computation.The reported overhead was 2.7% on SciFact and 1.75% on NFCorpus.Source: §I.3 · p. 26

Related papers

Loading 2406.05085v6…