PaperlayerView exact revision on arXiv

Source-linked AI summary

xRAG: Extreme Context Compression for Retrieval-augmented Generation with One Token

Xin Cheng, Xun Wang, Xingxing Zhang, Tao Ge, Si-Qing Chen, Furu Wei, Huishuai Zhang, Dongyan Zhao

arXiv:2405.13792v2cs.CLcs.AIcs.IR

TL;DR

Retrieval-augmented language models benefit from retrieved knowledge, but long documents can raise inference costs and exceed context limits. xRAG projects dense-retrieval document embeddings into the LLM representation space through a modality bridge, with the retriever and LLM frozen. The paper reports improvements across six knowledge-intensive tasks, comparable performance to uncompressed models on several benchmarks, and lower FLOPs.

  • Problem

    Retrieved documents can raise inference costs and exceed LLM context limits, while existing compression approaches face memory or compression-rate tradeoffs.

    Source: §1 · pp. 1–2

  • Method

    xRAG uses a modality bridge to project dense-retrieval document embeddings into the LLM representation space, keeping the retriever and LLM frozen.

    Source: §1 · p. 2

  • Results

    Over 10% improvement across six knowledge-intensive tasks is reported; xRAG is comparable to uncompressed models on several benchmarks and reduces total FLOPs by a factor of 3.53.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The trainable bridge allows reuse of pre-constructed document embeddings and preserves the plug-and-play nature of retrieval augmentation.

    Source: §1 · p. 2

    The study considers single dense-vector retrieval, not sparse or multi-vector retrieval, and reports a considerable performance gap versus RAG on reasoning tasks such as HotpotQA and FactKG.

    Source: G Limitations · p. 21

Abstract

from arXiv · show

This paper introduces xRAG, an innovative context compression method tailored for retrieval-augmented generation. xRAG reinterprets document embeddings in dense retrieval--traditionally used solely for retrieval--as features from the retrieval modality. By employing a modality fusion methodology, xRAG seamlessly integrates these embeddings into the language model representation space, effectively eliminating the need for their textual counterparts and achieving an extreme compression rate. In xRAG, the only trainable component is the modality bridge, while both the retriever and the language model remain frozen. This design choice allows for the reuse of offline-constructed document embeddings and preserves the plug-and-play nature of retrieval augmentation. Experimental results demonstrate that xRAG achieves an average improvement of over 10% across six knowledge-intensive tasks, adaptable to various language model backbones, ranging from a dense 7B model to an 8x7B Mixture of Experts configuration. xRAG not only significantly outperforms previous context compression methods but also matches the performance of uncompressed models on several datasets, while reducing overall FLOPs by a factor of 3.53. Our work pioneers new directions in retrieval-augmented generation from the perspective of multimodality fusion, and we hope it lays the foundation for future efficient and scalable retrieval-augmented systems

1 Introduction

xRAG addresses the cost of long retrieval-augmented prompts by integrating document embeddings into the language model with one token. The paper reports gains across knowledge-intensive tasks and reduced computation.

1 Introduction

  • Retrieved documents can expand the original query by more than tenfold, increasing inference costs and potentially exceeding LLM context limits.Prior compression methods either require substantial memory for stored LLM activations or have relatively low compression rates; the paper notes that dense-retrieval embeddings already distill document content.Source: §1 · pp. 1–2
  • xRAG projects dense-retrieval document embeddings into the LLM representation space through a modality bridge, leaving the retriever and LLM frozen.The bridge is the only trainable component, allowing reuse of pre-constructed document embeddings and preserving plug-and-play retrieval augmentation.Source: §1 · p. 2
  • Over 10% improvement across six knowledge-intensive tasks is reported with one document token, alongside 3.53× lower total FLOPs than uncompressed RAG.xRAG also delivers results comparable to uncompressed models on several benchmarks.Source: §1 · p. 2

2 Related Work

Related work spans retrieval-augmented language-model architectures and methods for compressing context. xRAG distinguishes itself by using modality fusion to integrate retrieval embeddings rather than relying on document text.

2 Related Work

  • Unlike standard RAG, which concatenates retrieved documents with the query, xRAG projects document embeddings into the LLM representation space using modality fusion.This reduces the added context to one token.Source: Figure 2 · p. 3
  • RAG research includes architectural changes, joint retriever–LM optimization, retriever alignment, and approaches to richer interaction.Other directions include iterative retrieval and generation and self-reflection.Source: §2 · p. 3
  • Context compression methods include soft memory-slot approaches [58, 19, 14, 47], token-level prompt compression [28, 29, 62, 82], and RAG-specific methods.A concurrent approach uses passage embeddings for listwise reranking.Source: §2 · p. 3

3 Methods

xRAG retrieves document embeddings and projects them into the language model’s representation space, replacing document text with a single feature token. It trains this bridge in two stages while keeping the language model and retriever frozen.

3.1 xRAG Architecture

  • xRAG replaces the document text in the LLM input with W(E), reducing input length from |D| + |q| to 1 + |q|.The query remains represented through the LLM embedding layer; E is the retrieved document’s dense representation.Source: §3 · p. 3, §3.1 · p. 4

3.2 Paraphrase Pretraining

  • Paraphrase pretraining teaches the model to recover document D from W(E), using an instruction and autoregressive next-token prediction.The objective aligns the retrieval feature with its corresponding document.Source: §3.2 · p. 4

3.3 Context-aware Instruction Tuning

  • Context-aware instruction tuning combines correct-answer language modeling conditioned on dense context features with self-distillation from the full-context RAG model.The final loss combines the language-modeling and KL-divergence losses using a hyperparameter.Source: §3.3 · pp. 4–5

3.4 Design Principle

  • A two-layer MLP projector is xRAG’s only trainable component, accounting for 0.46% of Mistral-7b parameters and 0.07% of Mixtral-8x7b parameters.xRAG reuses offline-constructed document embeddings and adds no memory overhead for storing LLM activations.Source: §3.4 · p. 5

4 Experimental Setup

The evaluation covers six knowledge-intensive tasks, uses Mistral-7b and Mixtral-8x7b with a Wikipedia retrieval corpus, and compares xRAG with non-retrieval, retrieval, and compression baselines.

4.1 Evaluation Dataset

  • Six knowledge-intensive datasets cover open-domain, multihop, and long-form question answering, plus fact-checking.Natural Questions, TriviaQA, and Web Questions are open-domain QA; HotpotQA is multihop QA; TruthfulQA is long-form QA; FactKG is fact-checking.Source: §4.1 · p. 5

4.2 Implementation Details

  • Experiments use Mistral-7b and Mixtral-8x7b with top-1 retrieval from approximately 37 million Wikipedia passages averaging 180 tokens.The corpus is the December 2021 Wikipedia dump, and SFR is the default retriever.Source: §4.2 · p. 6

4.3 Baselines

  • Baselines include non-retrieval and retrieval-augmented LLMs, LLMLingua, and TF-IDF compression; the LLMs remain frozen and retrieved documents are shared across methods.The table reports results on six downstream tasks and labels relative improvement over the non-retrieval setting.Source: §4.3 · p. 6, Table 1 · p. 6

5 Experimental Results

Across Mistral-7b and Mixtral-8x7b, retrieval augmentation significantly improves results, with xRAG compressing documents to one token while approaching or exceeding uncompressed performance; it also reduces computation versus RAG.

5.1 Knowledge Intensive Tasks

  • Across Mistral-7b and Mixtral-8x7b, retrieval augmentation significantly improves results (p-value < 0.05), with more modest gains for larger configurations.Source: §5.1 · pp. 6–7
  • 175.1 to one token: xRAG matches or exceeds uncompressed performance in some instances, with 16.6% versus 19.4% for Mistral-7b and 9.7% versus 8.0% for Mixtral-8x7b.The comparisons are xRAG versus the uncompressed model.Source: §5.1 · p. 7

5.2 Computational Efficiency

  • x1.64 increase in CUDA Time efficiency and x3.53 reduction in GFLOPs: xRAG outpaced RAG in profiling across four datasets.Measurements used Mistral-7b in bfloat16 inference mode with fixed batch size.Source: Table 2 · p. 7, §5.2 · p. 7

6 Analysis

The analysis finds xRAG resilient to noisy or misleading retrieval while documenting a boost-rate trade-off, training and data-blend findings, and a misleading-document case.

6.1 Evaluation Beyond the Overall Score

  • 75.2% average resilience for RAG means retrieval can harm about one-quarter of previously correct answers; xRAG is robust across datasets but trails RAG in boost rate, especially in multi-hop reasoning.Source: §6.1 · pp. 7–8, Figure 4 · p. 8

6.2 What makes xRAG effective?

  • For Mistral-7b, pretraining and finetuning matter equally; for Mixtral-8x7b, pretraining matters less, while self-distillation chiefly improves resilience over NLL finetuning.NLL optimization tends to increase reliance on context, making the system more vulnerable when relevant documents are not retrieved.Source: Table 3 · p. 8, §6.2 · pp. 8–9

6.3 Case Study

  • When retrieval misleads, standard RAG follows the document into a factually false answer, whereas xRAG relies on the LLM's internal knowledge.Source: Figure 5 · p. 9, §6.3 · pp. 9–10

6.2 What makes xRAG effective?

  • Reading-comprehension tuning most improves resilience and boost rates, while summarization-only tuning transfers to unseen QA, pointing to better use of projected context rather than task-specific knowledge.Source: Table 4 · p. 9, §6.2 · p. 9

7 Conclusion

xRAG is a context-compression method for retrieval-augmented generation that uses modality fusion. The conclusion reports faster inference than RAG with comparable performance on knowledge-intensive tasks.

7 Conclusion

  • xRAG can be significantly faster than RAG while maintaining comparable performance on knowledge-intensive tasks.The paper identifies reasoning over embeddings, handling multiple documents, and combining with multi-vector retrieval as areas for future improvement.Source: §7 · p. 10

A Comparison between different Context Compression Models

B Instructions for Paraphrase Pretraining

C Details for Instruction Tuning Dataset

The instruction-tuning dataset is assembled from raw data across categories using FLAN templates, with retrieval used to supply context for QA tasks lacking explicit context.

C Details for Instruction Tuning Dataset

D Implementation Details

The implementation uses instruction-tuned Mistral-7b and Mixtral-8x7b models, and preconstructed retrieval data to support efficient iteration. The embedding-model analysis finds that stronger sentence representation tends to improve downstream performance, with Dragon outperforming BGE-Large despite being smaller.

D Implementation Details

  • The evaluations use instruction-tuned Mistral-7b and Mixtral-8x7b, with Mixtral approximately 6.5 times larger and using a mixture-of-experts rather than dense architecture.Source: D Implementation Details · p. 18
  • To avoid on-the-fly retrieval, the experiments use a preconstructed ColBERT-v2 index for top-1 document retrieval, then re-encode documents with the embedding model of interest.The authors state this setup supports faster data-centric iteration.Source: D Implementation Details · p. 18
  • Tables 9 and 10 list hyperparameters for paraphrase pretraining and context-aware instruction tuning.Source: D Implementation Details · p. 18, Table 9 · p. 18, Table 10 · p. 18

E About different Embedding Models

  • Dragon, despite being BERT-base-sized, outperforms BGE-Large, which is twice its size; stronger sentence representation generally tends to improve downstream performance.The comparison averages performance over four QA datasets while holding the top-1 retrieved document fixed across embedding models.Source: E About different Embedding Models · p. 19, Table 11 · p. 19

F Analysis on Mistral-7b

The Mistral-7b analysis presents resilience and boost rates for four QA datasets and reports a pattern matching the Mixtral-8x7b model.

F Analysis on Mistral-7b

G Limitations

The paper identifies limits in xRAG’s retrieval coverage, reasoning performance, and document setting. These boundaries also define directions the authors consider promising.

G Limitations

  • xRAG is evaluated only with single dense-vector retrieval, excluding sparse methods such as BM25 and multi-vector methods such as ColBERT.The authors suggest combining these retrieval methods with xRAG as a promising direction.Source: G Limitations · p. 21
  • xRAG lags RAG considerably on reasoning tasks such as HotpotQA and FactKG, and the authors suggest its training data lacks reasoning-relevant examples.Improving xRAG’s reasoning ability remains future work.Source: G Limitations · p. 21
  • The evaluation considers only Top-1 retrieval, while the authors see potential for multi-document xRAG, whose input length scales by a factor of 1.The passage contrasts this with RAG, whose input length scales by the document-length factor.Source: G Limitations · p. 21

H More interesting cases

The examples show xRAG handling some long-document and retrieval-error cases that mislead RAG, while reasoning cases remain challenging. The checklist points to the paper’s reproducibility and limitations discussions.

H More interesting cases

2. Limitations

  • The checklist’s limitations response refers readers to Section 5.1 and Appendix G.Source: §2 · p. 24

NeurIPS Paper Checklist

  • The checklist marks the experimental results as reproducible and points to Appendix D for full training-process details.It also states that the models and data used are publicly available.Source: §2 · pp. 25–26

Related papers

Loading 2405.13792v2…