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

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

Paper

References (70)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer 2bxX4/10 · confidence 5/52024-07-07

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.

Strengths

- The compression rate from the long context to one token is high - The only trainable component constitutes less than 0.1% of the LLM’s parameters, which can serve as an efficient plug-in - xRAG can match or surpass the original RAG with only one token for one document, and the analysis of the Resilience Rate and the Boost Rate is reasonable and interesting - In general, the paper is well written

Weaknesses

- 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

- 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

Rating

4

Confidence

5

Soundness

2

Presentation

3

Contribution

2

Limitations

Yes.

Reviewer 2bxX2024-08-11

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.

Authorsrebuttal2024-08-11

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.

Reviewer MKxH3/10 · confidence 5/52024-07-11

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.

Strengths

1. This paper introduces xRAG, a novel context compression method that effectively merges document embeddings into the language model’s representation space through modality fusion, which is an attractive topic. 2. The xRAG achieves significant context compression compared with baselines. 3. The plug-and-play nature makes it easy to be applied with other backbone models.

Weaknesses

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

See weakness.

Rating

3

Confidence

5

Soundness

2

Presentation

3

Contribution

2

Limitations

Yes

Reviewer jbwR4/10 · confidence 4/52024-07-11

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.

Strengths

1. Using modality fusion bridge to connect the embeddings from dense retrieval and LLMs for retrieval-augmented generation. 2. Extremely compress the input tokens from RAG to only one token. 3. Comparable performance with lower FLOPs.

Weaknesses

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

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?

Rating

4

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

Authors have discussed the limitations of this paper in Section G of Appendix.

Reviewer H8JS7/10 · confidence 4/52024-07-14

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.

Strengths

- The proposed xRAG method treats document embedding as a modality, and shows strong potential for efficient RAG systems. Specifically, the proposed method uses lightweight training strategy that consists of paraphrase pretraining and context-aware instruction tuning. - xRAG demonstrates strong empirical results: it can achieve similar performance with traditional RAG methods on various QA datasets, and has the best performance among compression-based RAG methods. However, it still has limitations in tasks that require multi-step reasoning. - The paper provides in-depth analysis of the method's behavior, including studies on robustness, effectiveness of different components, and failure cases. - The writing is very clear and easy to understand.

Weaknesses

- 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

Please refer to the previous section.

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

Yes, the paper addresses the limitations.

Authorsrebuttal2024-08-12

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

Authorsrebuttal2024-08-12

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

Reviewer MKxH2024-08-13

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.

Authorsrebuttal2024-08-13

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

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC