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

Authors: xin  cheng
xun  wang
xingxing  zhang
tao  ge
si-qing  chen
furu  wei
huishuai  zhang
dongyan  zhao
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 09:50:16

Abstract: 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.

Main paper: 
 introduction: 
 Retrieval-Augmented Language Models (RALMs) [42, 22, 9, 14, 70] have shown exceptional performance in a variety of knowledge-intensive tasks. By retrieving domain-specific, long-tailed, and up-to-date knowledge from a non-parametric datastore, RALMs significantly extend the boundaries of parametric Large Language Models (LLMs). However, the integration of entire documents into prompts can significantly increase inference costs and may surpass the context limit of LLMs [29, 79] . As illustrated in Figure 1 , while the inclusion of a relevant document enables the LLM to generate accurate responses, it does so at the expense of processing documents that expand the original query by more than tenfold.How might we mitigate the costs associated with extended context while maintaining the benefits of retrieval augmentation? Recent research interest has converged on a promising direction: Context Compression. This concept is pursued through two primary strategies: soft-prompting methods, such as Gist [59] , AutoCompressor [15] , and ICAE [20] , which compress the context into dense memory slots, and hard-prompting methods, such as LLMLingua [29] and RECOMP [79] , where compression is applied on the surface form. These approaches, however, either require significant memory for storing LLM activations (e.g., 1.05 MB per token as reported by [59] ) or suffer from relatively low compression rates. More critically, these methods overlook a crucial characteristic of RALMs: through large-scale contrastive learning with question-document pairs, modern dense retrieval systems already distill document content into a single high-dimensional embedding, and this embedding reveals (almost) as much information as the text [58, 44] .In this paper, we pioneer an innovative approach to retrieval augmentation and context compression through the lens of modality fusion. Drawing from multimodal research, where text-only language models are taught to "perceive" and "listen," a pretrained modality encoder like CLIP [64] is typically used to extract modality features. These features are then integrated into language models using a modality fusion bridge [45, 52, 34] . Building on the conceptual overlap between the retriever encoder and modality encoder, we introduce xRAG. This model redefines document embeddings from dense retrieval-traditionally solely for retrieval purposes-as retrieval modality features. xRAG employs a modality fusion methodology to seamlessly integrate these embeddings into the language model's representation space, thus obviating the need for textual counterparts and achieving significant context compression. In xRAG, the modality bridge is the only trainable component, while both the retriever and the LLM are kept frozen. This design decision facilitates the reuse of pre-constructed document embeddings and maintains the plug-and-play nature of retrieval augmentation-two essential factors for a functional RAG system.To verify the effectiveness and versatility of our framework, we conducted comprehensive experiments with different LLM backbones, ranging from a dense 7B model to an 8x7B Mixture of Experts model. Our results reveal that adding just one document token could lead to over a 10% improvement across six knowledge-intensive tasks, significantly surpassing previous compression methods. xRAG also delivers results comparable to uncompressed models on several benchmarks. This is remarkable considering that the only trainable component constitutes less than 0.1% of the LLM's parameters. In terms of efficiency, xRAG reduces total FLOPs by a factor of 3.53 compared to the uncompressed RAG model. We further provide detailed analyses of xRAG, examining various training strategies, data blends, and component selections for the retrieval system. We believe this research sets a strong foundation for the development of future efficient and scalable retrieval-augmented systems.We carefully ablate four optimization choices: pretraining, instruction tuning, and two optimization objectives-language modeling (nll) and self-distillation (self-kd). We also train a Mistral-7b with LoRA on our instruction tuning dataset to rule out the possibility that our improvement simply comes from tuning on more data. The outcomes are presented in Table 3 .Refer to the background document and answer the question:Background: The Old Man and the Key | " several ""washed-up"" celebrities, was composed by Alf Clausen and written by Vitti. Even though Vitti received sole credit for writing the song, parts of the lyrics were written by Simpsons writer Carolyn Omine, which were then revised by fellow writer Matt Selman. Selman was especially satisfied with the rhyme sung by Mr. T; he stated that whenever he feels down, he thinks of that rhyme and that it "boosts [him] up inside." The episode features American actress Olympia Dukakis as Grampa's love interest Zelda. American comedian Bill Saluga also makes an appearance, as his television character Ray J. Johnson. The owner of the minivan is portrayed by series regular Hank Azaria, who imitated the voice of Clark Gable for the character."Question: Which US comedian/actor is credited with saying 'A man is only as old as the woman he feels'?Output: Bill Saluga.Answer the question:Question: Which US comedian/actor is credited with saying 'A man is only as old as the woman he feels'?Output: Groucho Marx.Refer to the background document and answer the question:Question: Which US comedian/actor is credited with saying 'A man is only as old as the woman he feels'?Output: Groucho Marx. performance of our framework. In the case of Mistral-7b, pretraining and finetuning phases are of equal significance to the end results. However, for Mixtral-8x7b, the impact of pretraining is notably diminished, likely due to the larger model's enhanced capability to incorporate multi-modality information. Furthermore, we find that during finetuning, self-distillation is more important than language modeling. The primary advantage of self-distillation lies in bolstering the resilience rate of the xRAG system. Optimization with nll loss tends to cause an overreliance on context information, rendering the system more vulnerable when the retriever fails to fetch relevant documents. To explore the effects of different data blends, we instruction-tune three xRAG model variants, each using data from these distinct categories. The results are shown in Table4. Our analysis reveals that among the dataset blends, reading comprehension data most significantly enhances the xRAG model's performance, as evidenced by both high resilience and boost rates. Intriguingly, when tuned solely with summarization data, xRAG still manages to deliver strong performance on QA datasets it has never been exposed to. This finding underscores that the advantages of instruction tuning for xRAG are not rooted in task-specific knowledge. Instead, they derive from the model's improved ability to utilize projected context information effectively. 
related work: 
 Retrieval-augmented Generation Equipping a parametric language model with a non-parametric datastore has proven effective for a range of NLP tasks, including language modeling [37, 57, 87] , open-domain question answering [25, 42, 70, 86] , domain adaptation [7] and machine translation [36, 12] , among others. Given the vast design space of this generation paradigm, numerous approaches with different focuses have been proposed. For instance, RETRO [9] and PlugLM [13] introduce architectural innovations for enhanced integration with the non-parametric datastore. REALM [22] pioneers an end-to-end approach for simultaneous optimization of the language model and retriever. REPLUG [70] and RA-DIT [51] improve retriever alignment using feedback from LLMs. DSP [38] and InteR [18] investigate complex interactions between the retriever and the language model. Selfmem [14] utilizes a reward model to refine retrieval and generation iteratively. Self-RAG [4] incorporates a self-reflection mechanism to enhance the quality and factuality of language model outputs. For a detailed overview, see [19, 5, 3] . Our contribution, xRAG, stands out by implementing a modality fusion approach to retrieval augmentation, creating an effective and efficient RAG system.Context Compression Context compression, aimed at reducing the input length for LLMs while retaining essential information, has recently attracted substantial interest [47] . Gist [59] achieves a compression rate of up to 26x by modifying the attention mask and caching soft gist token activations. ICAE [20] AutoCompressor [15] , and 500xCompressor [48] condense lengthy contexts into succinct, compact memory slots, which are directly utilizable by LLMs for diverse functions. LLMLingua [29, 30, 62] and CompAct [82] introduces a coarse-to-fine prompt compression technique based on perplexity scores and distilled token-level score. While these methods are generally applicable, others are tailored specifically for RAG systems, such as FilCo [76] and RECOMP [79] . A concurrent work directly employs passage embeddings for efficient listwise reranking [53] . For an in-depth comparison of these compression methods regarding memory efficiency, compression rates, and adaptability, refer to Appendix A. 
methods: 
 Problem Formulation In retrieval-augmented generation, a non-parametric datastoreD = {(E i , D i )} |D| i=1 consists of pairs where each D i represents a document chunk as a sequence of L i tokens D i = {d i 1 , . . . , d i Li }.Correspondingly, E i is the dense representation derived from a sentence embedding model SE θ (•) with input D i . For an input query q, its dense representation SE θ (q) is used to find the relevant documents by matching against the collection {E i } |D| i=1 with certain similarity search algorithm such as MIPS. After retrieval, the system selects a relevant pair (E, D) from D, concatenates the chosen document D with q, and processes the combined input with a language model F ϕ (D ⊕ q). Optionally, a context compression module C can be integrated to reduce the length of D from L to a more concise l, achieving a compression ratio of L l .Traditional methods for document compression typically focus on surface form of the document [29, 76, 79] . In contrast xRAG tackle the problem from a modality fusion view. Concretely, we introduce a modality projector W, which is trained to directly project the retrieval features E into the LLM representation space. Our proposed framework is visually contrasted with the traditional RAG system in Figure 2 . In the standard RAG, the input to the LLM comprises the embeddings Emb(D ⊕ q) of length |D| + |q|, where Emb signifies the embedding layer of the LLM. Conversely, with xRAG, the modified input is represented as W(E) ⊕ Emb(q), which yields a substantially reduced length of 1 + |q|. In this framework, the challenges come from the modality fusion: How can a text-only language model understand features from retrieval modality? To achieve this, we explore a two-stage training strategy: Paraphrase Pretraining followed by Context-aware Instruction Tuning.Similar to the pretraining strategies employed in vision-language models that use image-captioning data to align two modalities [52, 16, 54] , the primary objective of our paraphrase pretraining is to build a compatible representation between the extracted retrieval feature and the corresponding document. Illustrated in Figure 3 (a), for each pair (E, D) in a retrieval corpus D, we employ a natural language instruction X instruct to prompt the LLM to undertake a paraphrasing task (e.g. "[X] The above text could be paraphrased as: [D]", where [X] and [D] are placeholders for W(E) and document D) 2 . In this setup, the model learns to connect W(E) and D by recovering D on the condition of W(E) and the model is optimized by:EQUATIONwhere p ϕ is given by the softmax distribution of LLM F ϕ , and d <i denotes the document token before current prediction token d i , achieved by casual attention mask in auto-regressive LMs.After the pretraining phase, although the language model F ϕ has developed an internally compatible representation, it has never been explicitly trained to utilize these features for downstream tasks. Subsequently, we refine the optimization of on two directions:Optimization I: Language Modeling. Aligned with established instruction tuning methodologies [75, 24, 51, 43] , our objective is to finetune the model so that it generates the correct output when provided with a specific instruction, conditioned upon the given context information. Unlike traditional models that utilize the textual context X context , our method employs a dense feature E context to encapsulate the context information:EQUATIONOptimization II: Self-Distillation. The second trajectory of optimization aims to guide the xRAG in the effective utilization of contextual information, drawing from the principles of self-distillation [2, 71] and imitation learning [61, 23] . By considering the RAG model as a "teacher" and xRAG as a "student", we endeavor to distill the knowledge from RAG, thereby enabling xRAG to emulate the RAG model's proficiency in handling the full, uncompressed documents. This approach enhances xRAG's resilience in scenarios where it encounters noisy or irrelevant context that may not directly lead to the correct answer, detailedly discussed in § 6.1. Concretely, for a language model F ϕ using either X context or E context as the source of context, our objective is to minimize the divergence between the two resulting output distributions. This discrepancy is measured using the Kullback-Leibler (KL) divergence:EQUATIONHere X question is omitted for brevity and the final loss is the linear combination controlled by a hyperparameter: L nll + αL kl .In designing the projector W, our primary objective is to maintain the simplicity of the framework. We therefore opted for a two-layer MLP while other more sophisticated module such as Q-Former [45] could also be considered. Notice that the projector is the only trainable component, accounting for only 0.46% of the total parameters in the Mistral-7b model and 0.07% in the Mixtral-8x7b model. Such a design choice departs from previous studies that necessitated full-parameter tuning to adapt LLMs for compressed contexts [59, 76, 15] . We believe this approach will likely be more accessible and practical because, fundamentally, the RAG itself functions as a plug-and-play module for LLMs, and so should its compressed version. This design also avoid the risk of compromising other core capabilities of LLM during full-parameter tuning, as observed in [55, 56] .Moreover, in contrast to other compression methods that necessitate storing LLM activations for each compressed token [59, 20, 15] -an impractical strategy in the RAG setting, given the millions of documents involved-our method introduces no additional memory overhead. Instead, it leverages offline-constructed document embeddings, originally designed for retrieval. To summarize, xRAG not only simplifies the integration process but also avoids unnecessary computational or memory expenses.4 Experimental SetupQuestion: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [NA] Justification: the paper does not include theoretical results. Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
evaluation dataset: 
 We evaluated the performance of xRAG primarily on knowledge-intensive tasks which encompass a range of challenges: (1) three Open Domain Question Answering datasets that address questions on a wide array of topics: Natural Questions [41] , TriviaQA [33] , and Web Questions [8] . (2) one Multihop Question Answering dataset, HotpotQA [81] , which necessitates multi-step reasoning to generate answers. (3) one Long-form Question Answering dataset, TruthfulQA [50] , that requires the generation of long-form and truthful responses. ( 4) one fact-checking dataset, FactKG [39] , which challenges the model to use complex reasoning to determine the factual accuracy of given claims.In line with the KILT [63] and GenRead [84] , we assessed three ODQA datasets and HotpotQA using the Exact Match (EM) metric, FactKG with Accuracy, and for the long-form QA, we used both the F1 score and Rouge-L (R-L) score. These tasks demand a broad spectrum of world knowledge and have been extensively explored in the retrieval-augmentation literature [42, 22, 35, 10, 25, 70] .To demonstrate the versatility of our framework, we choose two backbones, differing in scale and architecture: Mistral-7b [27] and Mixtral-8x7b [28] . For the retrieval corpus, we utilized the Wikipedia dump from December 2021, which was pre-processed into passages following the methodology described in [26] . This resulted in approximately 37 million passages, each averaging 180 tokens in length. Our default retrieval model is the SFR [1] , which, at the time of writing this paper, holds the leading position on the MTEB leaderboard [60] . We use top-1 ranked document for inclusion in our instruction-tuning dataset and for the evaluation of downstream tasks. More details are provided in Appendix D. Code is available at: https://github.com/Hannibal046/xRAG.In determining appropriate baselines for comparison, we adhered to a fundamental principle: the selected compression methods must support the general plug-and-play capability of retrieval augmentation. This entails that they should function effectively without the need for dataset-specific tuning [79, 76] , or any alteration to the parameters of LLMs [76, 59] . Furthermore, given the extensive volume of the retrieval corpus, it is essential that these compression methods demonstrate memory efficiency, specifically by not requiring the storage of LLM activations for each individual token [59, 20, 15] . With these criteria in mind, our evaluation compares xRAG to the following baselines: (I) Primarily, we consider two variants of LLMs: one that operates without retrieval augmentation and another that includes it. These serve as the lower and upper performance bounds for our study of compression techniques, respectively. (II) Additionally, our comparisons extend to LLMLingua [29] , a plug-and-play approach for context compression. (III) Taking inspiration from [59] , we incorporate a method of discrete compression using TF-IDF. This approach yields compression rates comparable to those achieved by xRAG and serves as the lower bound for discrete compression. 5 Experimental ResultsQuestion: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: in Appendix D we provide full details about our training process Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
knowledge intensive tasks: 
 In Table 1 , we present our main results. Across both Mistral-7b and Mixtral-8x7b configurations, we observe a consistent and significant improvement when retrieval augmentation is applied (p-value < 0.05), although the gains are more modest for the larger model configurations. This trend aligns with observations reported by [70, 51] . Further analysis on the efficacy of various compression techniques reveals that xRAG outperforms other approaches by a large margin. Remarkably, xRAG not only reduces the token count drastically-from 175.1 to a single token-but also maintains robust performance levels. In some instances, xRAG's performance is comparable to, or even exceeds, that of the uncompressed models. Specifically, in the Mistral-7b configuration, xRAG achieves nearly the same performance improvement as the uncompressed model (16.6% compared to 19.4%), and in the Mixtral-8x7b configuration, it surpasses the uncompressed model (9.7% compared to 8.0%). One possible reason lies in the vulnerability of current RAG system when the irrelevant or misleading documents are presented, a topic detailed discussed in §6.1. We also observe that xRAG performs well in tasks that require document understanding, such as TriviaQA. However, in tasks that demand reasoning over document, like HotpotQA and FactKG, xRAG lags behind by a considerable gap. In this section, we conduct a thorough assessment of our framework's computational efficiency and memory management. To rigorously evaluate our model, we employed Torch Profiler3 to measure the CUDA Time (milliseconds) and Giga FLOPs of both the RAG and xRAG models across four real-world datasets. In these evaluations, the Mistral-7b, operating in bfloat16 inference mode, served as the base LLM. CUDA Time and GFLOPs were calculated on an average per batch basis with a fixed batch size, and GFLOPs were normalized by the number of generated tokens. These experiments were performed on the same computational hardware, specifically an Nvidia A100 and an AMD EPYC 7V12 64-Core Processor. As depicted in Although retrieval augmentation generally boosts performance as shown by aggregate metrics, it may not uniformly benefit all instances. In certain cases, the retrieval system might provide irrelevant or even misleading information, leading to incorrect answers that were previously correct [83, 76, 5] .To enable a more fine-grained evaluation, we introduce two novel metrics: the Resilience Rate and the Boost Rate. The resilience rate quantifies the percentage of instances in which the system's responses remain correct both before and after retrieval augmentation, highlighting the system's stability and robustness. Conversely, the boost rate measures the percentage of instances that were initially answered incorrectly but were rectified following the introduction of a retrieved document, thereby assessing the efficacy of retrieval augmentation. An ideal RAG system should have both high resilience rate and boost rate.In Figure 4 , we display these metrics for the uncompressed RAG and two compression methods: LLMLingua and xRAG. Surprisingly, although retrieval augmentation generally enhances performance, the resilience rate for RAG averages only 75.2%, indicating that retrieval can adversely affect about one-quarter of previously correct answers. In contrast, xRAG demonstrates considerable robustness across all evaluated datasets. This robustness largely stems from xRAG's ability to maintain an unbiased stance toward the internal knowledge representation of the LLM, especially when confronted with noisy retrieval content. Similar trends are noted in [51, 55] , where search-augmented instruction learning is shown to bolster the robustness of language models. However, xRAG still lags behind RAG in boost rate, particularly in multi-hop reasoning tasks. It is crucial to note that a high resilience rate does not necessarily mean that the LLM disregards the provided information, which could potentially lead to a reduced boost rate. A comparative analysis with LLMLingua indicates that xRAG is not only more robust but also more effective. This section delves into a thorough evaluation of various elements that contribute to xRAG's overall performance, focusing on its training strategy, the blend of datasets used and the effect of different embedding models. Due to the space limit, we present the last factor in Appendix E.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: all the models and data we use are publicly available and we carefully cite each paper Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
case study: 
 In Figure 5 , we show one interesting case about the robustness of xRAG. When retrieval system provide misleading content, standard RAG would overly rely on the document and generate answer that are faithful to the document while not factually true. Our xRAG model opt to rely on the internal knowledge of LLM and being robust to the misleading content. In Appendix H, we include more cases about xRAG including several error analysis.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: Section 4.1Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
conclusion: 
 In this work, we present xRAG, an innovative context compression method tailored for retrievalaugmented generation. For knowledge-intensive tasks, xRAG can be significantly faster than RAG while maintaining comparable performance. We are excited about the future of this modality-based retrieval-augmented system and plan to further improve its performance in the areas of reasoning over embedding, handling multiple documents, and combining with multi-vector retrieval.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: in Section 5.1Guidelines:• The answer NA means that the paper does not include experiments. Justification: reviewed and confirmed Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
acknowledgement: 
 We would like to express our sincere gratitude to the anonymous reviewers for their thorough review, insightful comments, and constructive suggestions, which have significantly improved the quality of this manuscript. This work paper is supported (in part) by the State Key Laboratory of General Artificial Intelligence. 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: in section 5.2Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: all the data and model we use is publicly available Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [NA] Justification: n/a Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: In Table 5 , we present a detailed comparison of various context compression models, emphasizing their real-world applicability. This comparison focuses on two key aspects: (1) Plug-and-play capability, which assesses whether dataset-specific tuning is necessary for new, unseen data; (2) Memory efficiency, which evaluates if additional memory space is required to store the compressed information, such as high-dimensional vectors typically used in soft prompting methods. The list of instructions for Paraphrase Pretraining is shown in Table 6 . They present the same meaning with natural language variance.• "Background: After collecting the raw data from different categories, we use templates 4 from FLAN [77] to construct instruction tuning dataset. In Table 7 we list an overview and in Table 8 we list the detailed information for each subtask of our dataset. For QA tasks that lack an explicit context, we perform a retrieval operation within the corpus D to identify the most relevant document to serve as context. This approach is akin to the retrieval-augmented instruction tuning depicted in [55, 51] . For the language models we use, Mixtral-8x7b is approximately 6.5 times larger in scale compared to Mistral-7b and features a divergent architectural approach-specifically, a dense versus mixture-ofexperts design. For our assessments, we employed the instruction-tuned variants of these models.Owing to efficiency constraints, we opted not to perform on-the-fly retrieval. Instead, we preconstructed a retrieval index using the efficient and robust multi-vector retriever, ColBERT-v2 [68] , from which we retrieved the top-1 ranked document for inclusion in our instruction-tuning dataset and for the evaluation of downstream tasks. Subsequently, we re-encoded these documents using the embedding model of interest (e.g., SFR). This strategy allows us to iterate data-centric experiments quickly. All experiments are conducted on the a setup of 8xNvidia A100 GPUs.In Table 9 In our primary experiments, we use the SFR model as our default sentence embedding model. This section delves into the effects of different embedding models. We examine four universal text embedding models: E5-Mistral and E5-Large [73] alongside BGE-Large and BGE-Base [78] . Additionally, we assess two retrieval-specific models: Dragon [49] and DPR [35] . The configurations of the different retrievers and their MTEB scores 6 -a metric indicating their general sentence representation capability-are listed in Table 11 . To isolate the impact of potentially different retrieved documents, we ensure that all models utilize the same top-1 document. The performance is averaged over four question answering datasets. A general pattern is that embedding models with stronger sentence representation capabilities tend to further enhance the downstream performance. Remarkably, the Dragon model, despite being a BERT-base-sized retrieval-specific model, outperforms general text embedding models that are twice its size (BGE-Large). In Figure 6 We discuss the limitations of our framework as follows:• In this work, we only consider the most commonly used retrieval system-single dense vector retrieval, while sparse retrieval methods such as BM25 or multi-vector retrieval methods like ColBERT are not included. We believe that combining these methods would be a promising direction for xRAG, as sparse vectors could complement dense vectors, and multi-vector retrieval would provide xRAG with more flexibility by not condensing all information into one token. • Currently, xRAG delivers decent performance when a relevant document is fetched; however, it lags behind RAG by a considerable margin in tasks that require reasoning (such as HotpotQA and FactKG). One possible reason is that during the training phase of xRAG, reasoning-relevant data is not provided. How to make xRAG a better reasoner remains our future work. • We only consider the Top-1 retrieval setting, while ensembling multiple relevant documents has been shown to be effective for RAG systems due to the complementary information contained in Top-K documents. We believe there is potential advantage for xRAG to scale to multi-document settings, as the input length of xRAG for multi-documents scales by a factor of 1, while for RAG, it scales by the document length factor.• In Figure 7 , we report a failure case of xRAG. In this case, retrieval alone is not enough to derive the final answer and the LLM is required to perform reasoning over retrieved document (the listed universities are all located in Switzerland). • An interesting example is shown in Figure 8 , when the retrieved document is a list of a characters in the book Discworld, the RAG model would respond with a fictional character, while xRAG generate the right answer by focusing on the relevant part of the document. • In Figure 9 , even when the retrieved document is relevant, RAG would still hallucinate while xRAG could generate the right answer based on the document. • In Figure 10 , the retriever mistakenly fetch the wrong document (Phantom of the Opera of interest is a music rather than a file) and RAG would be misled while xRAG remain robust to generate the correct answer.Refer to the background document and answer the question: Output: I don't know.Background: Philip Kraft | Philip Kraft (Refer to the background document and answer the question:Question: Philip Kraft lectured at the largest university in what European country?Output: Germany.Figure 7 : Failure case of xRAG when reasoning is required to derive the final answer.Refer to the background document and answer the question:Background: List of Discworld characters | " Meaning ""Two-Fingered"" in Ephebian, Didactylos is a philosopher, based on Diogenes of Sinope, who comes into the stories in Small Gods. He lives in a barrel inside the wall of the palace of the Tyrant in Ephebe, crafting bespoke philosophies, axioms or aphorisms for scraps. Having apprenticed his nephew Urn as a philosopher, they frequently over the merits of natural philosophy, specifically Urn's development of steam power. Although one of the most popular philosophers of all time, Didactylos never earns the respect of his fellow philosophers, due to the fact that he thinks 'about the wrong things'. He has been pictured with a lantern though blind and is looking for an ""honest man"".Question: Which Greek philosopher was said to have lived inside a barrel?Output: Didactylos.Answer the question:Which Greek philosopher was said to have lived inside a barrel? Output: No one.Refer to the background document and answer the question:Question: Which Greek philosopher was said to have lived inside a barrel? Output: Diogenes of Sinope. Answer the question:Question: What mission was the fifth human space flight of 'Project Apollo', the third human voyage to the Moon and the first manned mission to land on the Moon?Output: Apollo 9.Refer to the background document and answer the question:Question: What mission was the fifth human space flight of 'Project Apollo', the third human voyage to the Moon and the first manned mission to land on the Moon?Output: Apollo 11.Figure 9 : xRAG correctly locates the relevant part in a long document while RAG would still hallucinate the wrong answer.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.To maintain diversity, we sample from an instruction pool, which could be found in Appendix B.https://pytorch.org/docs/stable/profiler.html#module-torch.profilerhttps://github.com/google-research/FLAN/blob/main/flan/templates.pyhttps://huggingface.co/spaces/mteb/leaderboard 


Decision: The paper presents a context compression approach for retrieval-augmented generation (RAG). The idea is to use document embeddings that were used for retrieval systems and train a light-weight adaptor that can integrate information from such embedding into LLMs. This compresses in-context document significantly, resulting in much more efficient model (with smaller FLOPS). Overall, the idea is pretty simple but well-executed and shows strong experimental gains. 

The reviewers identified some limitations of the proposed approach - such as limited performance when handling multiple retrieved documents. While the approach shows efficiency gain in this setting, it comes with moderate performance drops. The new results from the discussion period should be added to the revised paper. With the advance of the long context model, extending this work to compressing long document setting will be necessary for future impact.

Reviews: review: paper_summary: paper_summary: This paper presents xRAG, a context compression method for Retrieval-Augmented Generation (RAG). Their key idea is to treat document embeddings from dense retrieval as features from a retrieval modality, which allows compressing retrieved documents into a single token. Experiments show that their method can achieve extreme compression while maintaining performance comparable with traditional RAG.
strength_weakness: strength_weakness: - RAG is proposed to reduce the hallucination of language model generation, by letting them refer to the original context. It is still unclear why xRAG that compresses a document into a single token could still leave the factual knowledge intact, and the authors should discuss more on this.
- As the limitation paragraph mentions, the paper does not consider the case of retrieving multiple documents. Though the proposed method works for one-document case, it may not generalize to multiple-document case and may affect the training of the compression projector.
questions: questions: Please refer to the previous section.
limitations: limitations: Yes, the paper addresses the limitations.
paper_summary: paper_summary: This paper proposes xRAG, a context compression method designed specifically for retrieval-augmented generation. xRAG redefines the use of document embeddings in dense retrieval by integrating them as features from the retrieval modality.  It achieves an extreme compression rate to only one token. The authors perform experiments to  demonstrate the effectiveness of the proposed method.
strength_weakness: strength_weakness: 1. The proposed method relies heavily on the performance of dense retrieval models. The dense retrieval models have poor generalization for out-of-domain [1], which does not match the general ability of LLMs, but the experiment in this paper is only conducted on wiki-based knowledge bases of Q&A tasks.

2. The specific details of selection and training method of dense retrieval model in xRAG are needed.

3. How xRAG performs on long-context RAG? Can only one token represents the semantics of the retrieved long-documents?


[1] Back to Basics: A Simple Recipe for Improving Out-of-Domain Retrieval in Dense Encoders
questions: questions: 1. How xRAG performs on long-context RAG? Can only one token represents the semantics of the retrieved long-documents?

2. How xRAG selects or trains the dense retrieval model?

3 How xRAG performs on the domain that shifts from the training domain of dense retrieval model?
limitations: limitations: Authors have discussed the limitations of this paper in Section G of Appendix.
paper_summary: paper_summary: The paper on xRAG presents an innovative approach to context compression in retrieval-augmented generation, achieving significant efficiency gains while maintaining performance. The method's compatibility with various language models and preservation of the plug-and-play nature are notable strengths. However, the added complexity and dependency on high-quality embeddings might pose challenges. The whole paper is clear and easy to understand.
strength_weakness: strength_weakness: 1. The introduction of a modality bridge and the requirement for a pretrained modality encoder adds complexity to the system.
2. Will the bias in projector tuning cause hallucinations?
3. The performance of xRAG is likely highly dependent on the quality of the dense document embeddings.
4. The compression method, while efficient, might result in the loss of information present in the original documents. Do we need such extreme compression since most of the queries are not too long? Or how to balance the compressed tokens with the performance?
5. The concepts are so confused. Sometimes, the author uses the "modality bridge" and also uses the "projector". I think they are the same module right?
6. The paper mentions using a two-layer MLP as the projector but doesn't explore how different projector architectures might impact performance.
7. The paper suggests that xRAG is more robust than RAG when dealing with irrelevant documents but doesn't provide a clear mechanism for how xRAG achieves this robustness.
8. How will the hyper parameter alpha influence the final results?
9. How to ensure the compressing process filters the irrelevant information rather than the crucial ones?
10. Will sensitive information like numbers or times be kept by the projector after compressing?
questions: questions: See weakness.
limitations: limitations: Yes
paper_summary: paper_summary: This paper proposes xRAG, a method to map the retriever's embeddings into LM's representation space with one token. While both the retriever and LM are fixed, xRAG trains a simple projector to adapt to them by paraphrase pretraining and context-aware instruction tuning. xRAG can significantly outperform non-retrieval and previous context compression methods and matches the performance of the uncompressed one, saving the RAG FLOPs by more than 3 times.
strength_weakness: strength_weakness: - The main experiment in this paper is conducted only on one retrieved document, whose average length is just 175, but I think this is a **very short context** for modern LLMs (and thus the inference time is not a bottleneck). I think proving the effectiveness of xRAG on a real-long context (more than thousands of tokens) of multiple documents is very important to see this method's real-world value

- It is essential to consider the total training FLOPs of the projector. Although we don't need the actual optimization of LM and retriever, we also need the forward process of these models in the projector's training so that **the cost can be significant**, especially when the amount of total two-stage training samples is very, very large (near 3M)

- The projector is **tightly coupled** with the retriever and LM. It is hard to prove the generalization abilities of xRAG to other black-box production LLMs where the projector cannot be trained directly from the LM's feedback, which may limit its potential impact

- In Section 4.3, the requirements of the baselines include **no need for dataset-specific tuning**. However, how to explain that NQ, TriviaQA (two evaluation tasks) are included in the Context-aware Instruction Tuning data?

- Missing baselines of context compression methods other than LLMLingua, like Gist-COCO. Will adding **more compression tokens** benefit the performance?

[1]: Li, Xinze, et al. "Say more with less: Understanding prompt learning behaviors through gist compression." arXiv preprint arXiv:2402.16058 (2024).
questions: questions: - For plugging in multiple retrieved documents, should we devise another form of training task or simply adopt the same projector multiple times? This is an interesting discussion that I would like to see
- How simple can the projector be? I would like to see the performance of different design choices of the projector, such as Q-Former mentioned in the paper
- What is the data source for Paraphrase Pretraining? I cannot find it in the paper
limitations: limitations: Yes.

Comments: comment: Dear Reviewer,

Thank you for taking the time to engage in further discussion and for acknowledging that some of your concerns have been addressed. We would like to clarify the rationale behind compressing the context into a single token.

As demonstrated in our paper, our approach **not only achieves a higher compression rate but also results in better accuracy compared to existing context compression methods.** We believe that improving both performance and efficiency is crucial in advancing the state-of-the-art. 

Could you please elaborate on why a context compression method with better performance and efficiency is not preferred?

We appreciate your insights and look forward to your feedback.

Sincerely,

Authors
Thanks to the author for the responses and some of my concerns have been solved. The necessity of compressing the context into one token remains questionable and I will keep my score.
Dear Reviewers,

We hope this message finds you well. As the end of the rebuttal phase is fast approaching (Aug 13 11:59pm AoE), we wanted to kindly request your feedback on the amendments and clarifications provided in the rebuttal. Your insights and critiques greatly contribute to enhancing the quality of our work.

We are more than willing to provide further clarifications or engage in additional discussions if needed.

Thank you once again for your time and consideration.

Sincerely,

Authors
Dear Reviewer,

Thank you for your continued engagement and insightful question. We appreciate the opportunity to provide further clarification.

As requested, we have conducted additional experiments comparing xRAG and RAG performance with multiple chunks. The results are as follows:

|           | NQ   | TriviaQA | WebQ | HotpotQA | Doc Length|
| --------- | ---- | -------- | ---- | -------- | ----- |
| xRAG-top3 | 41.4 | 67.3     | 41.1 | 35.3     | 3     |
| RAG-top3  | 45.6 | 69.1     | 40.4 | 40.2     | 525.6 |

We want to emphasize that this version of xRAG, as described in our general response, **serves as a feasibility test** rather than a rigorous evaluation of xRAG's full potential with multiple chunks. This implementation modifies only a subset of data (summarization) in our context-aware instruction tuning while maintaining the original model architecture, paraphrase pretraining stage, and objective function. Given that the summarization data is typically divided into at most 3 chunks (with a maximum length of 512 tokens per sample and 180 tokens per chunk), an evaluation with top-10 chunks is beyond the scope of the current xRAG implementation.

While xRAG currently falls short of RAG in this multi-chunk setting, we insist **this does not diminish our core contribution**. The expansion to multiple chunks represents an incremental improvement (1 to N) of the existing framework. While valuable, we consider it less fundamental compared to our primary  0 to 1 innovation: efficient RAG with modality fusion.

We appreciate your thoughtful consideration and hope this additional information addresses your query satisfactorily. We remain committed to advancing this promising research direction and value the constructive feedback from the review process.


Best regards,

Authors
Thanks for your response. How about xRAG's performance compared with RAG (no compression) in top-3 and top-10 chunks? A partial evaluation should be fine. I will increase my score if this question is clear.
Dear Reviewer,

We sincerely appreciate your thoughtful feedback and welcome the opportunity to address the points you've raised.

- **Concern: The additional training cost seems non-trivial**

Similar to how LLaMA[1] "violates" the scaling law by prioritizing the inference budget, we believe that inference cost is of utmost importance for a context compression method, especially in real-world RAG scenarios involving millions of documents.  This principle guided xRAG's design: by reusing existing document embeddings and introducing only a two-layer MLP into the existing LLM, we've minimized additional complexity.

While our method prioritizes inference-time efficiency, it's worth noting that **the training cost is also minimal compared to existing compression methods such as ICAE[2] and LLMLingua[3].** The only trainable component in xRAG is the newly introduced MLP. We would greatly appreciate if you could point us towards any references to context compression methods with trivial training costs that we may have overlooked.

---

- **Concern: The efficiency gain in handling super-long contexts (beyond just 3 documents) is not included.**

RAG is widely recognized as a technique designed to alleviate the need for LLMs to process long contexts by chunking documents into smaller pieces and retrieving the most relevant ones. This is the typical operational model of modern RAG systems. As such, long-context processing and RAG represent two distinct paradigms for LLMs.

Research has shown that useful information generally appears within the top-k documents, and RAG performance tends to plateau as more documents are involved [4] [5]. Consequently, handling super-long contexts falls outside the scope of RAG and, by extension, beyond the scope of xRAG.

However, to address your concern about efficiency with more documents, we've conducted additional tests. If the "super-long" context you mentioned falls within the top-k chunks of RAG (where k is typically less than 10), here are the efficiency results for top-10 documents, following the benchmark setting outlined in Section 5.2 of our paper:



| Top-10 chunks | CUDA Time (s) | GFLOPs         | Peak Mem (GiB) |
| ------------- | ------------- | -------------- | -------------- |
| RAG           | 3.58s         | 10712.22       | 20.42          |
| xRAG          | 0.62s (x5.7)  | 529.33 (x20.2) | 13.84 (x1.4)   |

As demonstrated, xRAG maintains significant efficiency gains even when processing a larger number of documents.

We appreciate your insightful feedback and look forward to further discussion on these points.

---
[1] Touvron, Hugo, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave and Guillaume Lample. “LLaMA: Open and Efficient Foundation Language Models.” ArXiv abs/2302.13971 (2023): n. pag.

[2] Ge, Tao, Jing Hu, Xun Wang, Si-Qing Chen and Furu Wei. “In-context Autoencoder for Context Compression in a Large Language Model.” ArXiv abs/2307.06945 (2023): n. pag.

[3] Jiang, Huiqiang, Qianhui Wu, Chin-Yew Lin, Yuqing Yang and Lili Qiu. “LLMLingua: Compressing Prompts for Accelerated Inference of Large Language Models.” Conference on Empirical Methods in Natural Language Processing (2023).

[4] Lewis, Patrick, Ethan Perez, Aleksandara Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Kuttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, Sebastian Riedel and Douwe Kiela. “Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks.” ArXiv abs/2005.11401 (2020): n. pag.

[5] Shi, Weijia, Sewon Min, Michihiro Yasunaga, Minjoon Seo, Rich James, Mike Lewis, Luke Zettlemoyer and Wen-tau Yih. “REPLUG: Retrieval-Augmented Black-Box Language Models.” ArXiv abs/2301.12652 (2023): n. pag.
Thanks to the authors. I have read the rebuttal carefully. I decided to not change my scores since I feel the limitations of this work still outweigh reasons to accept. (1): The additional training cost seems not trivial; (2): The efficiency gain in super-long context (not only just 3 docs) is not included.