Source-linked AI summary

LongRAG: Enhancing Retrieval-Augmented Generation with Long-context LLMs

Ziyan Jiang, Xueguang Ma, Wenhu Chen

arXiv:2406.15319v3cs.CLcs.AI

TL;DR

Traditional RAG burdens retrieval with many short units that may lose context and underuse long-context readers. LongRAG uses long retrieval units, a few-unit retriever, and a long reader, achieving strong results across four QA datasets without training and supporting the reconsideration of retrieval granularity.

  • Problem

    Traditional RAG uses short retrieval units, creating a heavy retriever/light reader imbalance and potentially losing contextual information.

  • Method

    LongRAG forms 4K-token-scale retrieval units, retrieves only a few coarse candidates, and uses a long-context LLM as the reader.

  • Results

    LongRAG boosts performance across NQ, HotpotQA, Qasper, and MultiFieldQA-en without additional training.

  • Takeaways & Limitations

    LongRAG suggests modern RAG systems can exploit long-context LLMs by reconsidering retrieval-unit granularity.

  • Takeaways & Limitations

    Encoding a complete long retrieval unit is challenging, so LongRAG approximates its similarity using constituent chunks.

Abstract

from arXiv · show

In traditional RAG framework, the basic retrieval units are normally short. The common retrievers like DPR normally work with 100-word Wikipedia paragraphs. Such a design forces the retriever to search over a large corpus to find the `needle' unit. In contrast, the readers only need to generate answers from the short retrieved units. The imbalanced `heavy' retriever and `light' reader design can lead to sub-optimal performance. The loss of contextual information in the short, chunked units may increase the likelihood of introducing hard negatives during the retrieval stage. Additionally, the reader might not fully leverage the capabilities of recent advancements in LLMs. In order to alleviate the imbalance, we propose a new framework LongRAG, consisting of a `long retriever' and a `long reader'. In the two Wikipedia-based datasets, NQ and HotpotQA, LongRAG processes the entire Wikipedia corpus into 4K-token units by grouping related documents. By increasing the unit size, we significantly reduce the total number of units. This greatly reduces the burden on the retriever, resulting in strong retrieval performance with only a few (less than 8) top units. Without requiring any training, LongRAG achieves an EM of 62.7% on NQ and 64.3% on HotpotQA, which are on par with the (fully-trained) SoTA model. Furthermore, we test on two non-Wikipedia-based datasets, Qasper and MultiFieldQA-en. LongRAG processes each individual document as a single (long) unit rather than chunking them into smaller units. By doing so, we achieve an F1 score of 25.9% on Qasper and 57.5% on MultiFieldQA-en. Our study offers insights into the future roadmap for combining RAG with long-context LLMs.

1 Introduction

LongRAG revisits short-unit RAG because it burdens retrieval, loses context, and underuses long-context readers. It balances the system with long retrieval units, few-unit retrieval, and a long reader, improving performance across four QA datasets without training.

  • 1 Introduction: LongRAG addresses short-unit RAG’s imbalanced heavy-retriever/light-reader design, which can lose context and increase hard negatives.Recent long-context LLMs motivate shifting more workload from retrieval toward reading.
  • 1 Introduction: 4K-token retrieval units reduce corpus size and enable retrieval with more complete information.Units are formed from entire documents or grouped related documents.
  • 1 Introduction: Only a few top units are retrieved without reranking, reducing the likelihood that hard negatives confuse the reader.The reported range is 1 to 8 retrieval units across the four datasets.
  • 1 Introduction: LongRAG boosts overall performance on NQ, HotpotQA, Qasper, and MultiFieldQA-en without additional training.The introduction reports improvements across Wikipedia-based and non-Wikipedia-based QA settings.
  • 1 Introduction: With a 40K recall-token budget, short units reach 91% recall@200 but lower end performance, whereas long units improve or plateau as recalled units increase from 1 to 8.The contrast is attributed to the short-unit setting’s greater hard-negative impact.

2 LongRAG

LongRAG forms long retrieval units, retrieves a few coarse candidates, and gives their concatenation to a long-context reader. Its retriever uses vector similarity with an approximation for encoding long units, while larger units correspond to smaller k.

  • 2 LongRAG: LongRAG shifts retrieval from exact short-answer matching toward coarse relevant-context recall, leaving answer extraction to the reader.The reader receives only a few long units rather than hundreds of short ones.
  • 2 LongRAG: Documents longer than 4K tokens remain single units, while shorter documents are grouped with related documents into larger units.This construction produces long retrieval units from whole documents or document groups.
  • 2 LongRAG: Long units preserve document context and reduce hard negatives by limiting reader input to roughly 4–8 retrieved units.Traditional RAG may feed hundreds of short units into the reader.
  • 2 LongRAG: The retriever encodes questions and retrieval units as d-dimensional vectors and ranks units by their dot-product similarity.The top k units are concatenated into the long retrieval context.
  • 2 LongRAG: Because encoding an entire long unit is difficult, LongRAG approximates its score by maximizing similarity over constituent chunks.The paper evaluates 512-token, 4K-token, and whole-unit chunk granularities.
  • 2 LongRAG: Larger retrieval units require smaller k: NQ uses more than 100 passage units, about 10 document units, or typically 4–8 grouped-document units.The choice of k decreases as retrieval-unit size increases.

3 Experiments

LongRAG is evaluated across four question-answering datasets using long retrieval units, retrieval ablations, encoding comparisons, and reader-model comparisons. Long-context retrieval improves recall efficiency and supports competitive QA performance with fewer retrieved units.

  • 3.2 Retrieval Performance: Long retrieval units compress the NQ corpus from 22M to 600K units, raise top-1 answer recall from 52.24 to 71.69, and achieve comparable recall with 8 units instead of 100.The average long retrieval unit contains up to 6K tokens.
  • 3.2 Retrieval Performance: The proposed long-unit encoding approximation outperforms encoding the entire long context with an existing long embedding model.The approximation maximizes query similarity over chunks within each retrieval unit.
  • 3.3 Full QA Performance on Wikipedia-based Datasets: 62.7 exact match on NQ and 64.3 exact match on HotpotQA place LongRAG near the strongest fully supervised RAG baselines without fine-tuning.The comparisons use closed-book, fully supervised RAG, and no-fine-tuning RAG baselines.
  • 3.5 Ablation Studies: Figure 3 shows that retrieval beyond the turning point harms reader performance; on NQ, the point occurs between 100–200 passages, 5–10 documents, or 4–8 grouped documents.The suitable reader context is generally around 30K tokens, and long retrieval units outperform passage-level units.

4 Related Work

Related work improves retrieval, long-context processing, position encoding, and embedding-model context length for knowledge-intensive language-model applications.

  • Retriever-reader architectures retrieve corpus information for a language model to use as additional context in knowledge-intensive tasks.
  • Transformer computation grows quadratically with sequence length, motivating sliding windows, chunk segmentation, and FlashAttention for long inputs.
  • RoPE and AliBI position encodings, alongside position reorganization and interpolation, support efforts toward length extrapolation.
  • Long-context embedding models extend supported snippets from 512 tokens to 32k tokens using long-input pretraining or existing long-context language models.

5 Conclusion

LongRAG addresses retriever-reader imbalance by using 4K-token retrieval units with long retrieval and reading components, achieving strong retrieval with few units and no training.

  • LongRAG uses 4K-token retrieval units, a long retriever, and a long reader to reduce corpus size and retriever burden.The framework preserves document semantic integrity and minimizes hard-negative noise through coarse retrieval.
  • The framework achieves strong retrieval recall with only a few top units and demonstrates superior performance across four end-to-end question-answering tasks without training.
  • LongRAG is presented as a direction for modern RAG systems that combine retrieval with long-context language models.

A.1 Prompts Template for Long Context Reader

The long-context reader uses a two-stage prompting procedure: first generate a longer answer from retrieved context, then extract a concise answer using examples.

  • The first prompt turn concatenates the long retrieved context and question without in-context examples to generate a longer answer.The initial answer typically ranges from a few words to a few sentences.
  • The prompt template instructs the reader to answer directly and avoid outputting unrelated material.

LongRAG

LongRAG’s long-context reader uses a two-turn prompt design to generate a detailed answer from context and then extract a concise final answer with demonstrations.

  • LongRAG uses a two-turn approach in which the first turn generates a longer answer and the second extracts the exact short answer.The first turn omits in-context examples, while the second uses 8-shot examples for calibration.

A.2 Refined Metric

LongRAG refines Exact Match to recognize short predictions containing the ground truth or an alias, addressing cases where exact string equality is too strict.

  • A.2 Refined Metric: A prediction under five tokens counts as an exact match when the ground truth is a substring of it or vice versa.The paper manually verified that this captures aliases and other ground-truth forms.
  • A.2 Refined Metric: The refinement is motivated by long retrieved contexts, which increase the chance that an answer alias appears and can be extracted.
  • A.2 Refined Metric: Table 8 illustrates cases where LongRAG’s prediction differs textually from the ground truth but is judged correct under the refined metric.

A.3 Group Documents Algorithm

The grouping algorithm constructs long retrieval units by combining related documents under a token limit, using document connectivity and group size to guide merging.

  • A.3 Group Documents Algorithm: Document relatedness can be determined from hyperlinks, word frequency, or structural dataset information; the NQ and HotpotQA experiments use hyperlinks.
  • A.3 Group Documents Algorithm: The algorithm accepts a token limit, documents, related-document adjacency lists, and document degrees, then outputs a set of groups.
  • A.3 Group Documents Algorithm: Documents are processed from low to high degree, while related groups are sorted by size before being merged into a new group.
  • A.3 Group Documents Algorithm: After merging, the algorithm removes absorbed groups, adds the new group to the collection, and returns the resulting groups.

A.4 Dataset Examples

The paper provides examples drawn from the four datasets used in its experiments.

  • A.4 Dataset Examples: The dataset-example section presents examples from all four experimental datasets.
  • A.4 Dataset Examples: The examples are presented as illustrative material rather than as a separate experimental result.
  • A.4 Dataset Examples: Table 9 is designated as a collection of examples from the four datasets used in the experiments.
Loading 2406.15319v3…