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

Dear Reviewer,

We greatly appreciate your time and effort in reviewing our paper. Here, we would like to address your concerns point by point:

• About multiple documents:

Please refer to our general response about multiple documents in xRAG.

• The cost of training

The training cost of xRAG is minimal compared to full parameter tuning since the only trainable component is a two-layer MLPs. In our setup (8xA00), full parameter tuning of a 7B dense model with 1 million samples takes about 24 hours. In the same amount of time, xRAG (8x7b) can finish training on 2.9 million samples (paraphrase pre-training + context-aware instruction tuning). Moreover, by caching all the document embeddings, we can further speed up the training by 25%. We appreciate your suggestion and will include this information in the revised version of our paper.

• xRAG for black-box LLM

From the user side, it is true that xRAG can only be applied to models with open weights. However, with the increasing popularity of cloud vector databases, LLM API providers could still benefit from xRAG if users can provide the index of retrieved documents in the vector database.

• About dataset-specific tuning.

When we refer to "without dataset-specific tuning," we mean that we adopt a one-for-all evaluation setting where the evaluation results from various datasets come from the same model, as opposed to one-for-one evaluation which requires dataset-specific tuning. To exclude the impact of the corresponding training set (NQ and TQA), we refer to Table 4 of our paper, where similar results are obtained even with only reading comprehension data involved (NQ and TQA excluded). Additionally, the LoRA results in Table 3 confirm that the improvement of xRAG does not stem from task knowledge inherited during context-aware instruction tuning.

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

Thank you for pointing out this missing reference. Similar to Gist Token [1], Gist-COCO relies on gist representation, which might not be applicable in real-world RAG settings involving millions of documents. We will add this paper to the final version of xRAG in both the Related Work section and Appendix A. Regarding the addition of more compression tokens, we believe it is a balance between efficiency and efficacy. In Appendix G, we discuss that xRAG could be combined with a multi-vector or Matryoshka embedding system where the compression ratio is configurable. We believe this is an exciting direction for xRAG.

• 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

Please refer to our general response about multiple documents in xRAG.

• 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

As discussed in Section 3.4, the key design principle is the simplicity of the system. To our knowledge, the simplest form for a projector is an MLP architecture. We leave the exploration of more advanced and complex architectures for future work.

• What is the data source for Paraphrase Pretraining? I cannot find it in the paper

Apology for the confusion. The data source for Paraphrase Pretraining is the same Wikipedia dump used as the retrieval corpus. We will make this clear in the revised version of our paper.

[1] Mu, Jesse, Xiang Lisa Li and Noah D. Goodman. “Learning to Compress Prompts with Gist Tokens.” ArXiv abs/2304.08467 (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.

Dear Reviewer,

Appreciate your time and effort in reviewing our paper. We would like to address your concerns point by point:

• Added complexity to the system

In xRAG, we have emphasized throughout our paper that our key design principle is to add minimal complexity to the existing RAG system. In Section 3.4, we explicitly state that our approach does not require an additional pretrained modality encoder but instead utilizes existing document embeddings from the retrieval datastore. This design choice ensures zero additional space complexity, enabling practical scalability in real-world RAG systems handling millions of documents. To make our method plug-and-play, we freeze the LLM and retriever, and the only trainable component is the projector.

• Will the Bias in Projector Tuning Cause Hallucinations?

In Section 6.1, we introduce a metric called Resilience Rate, which measures how retrieval (RAG or xRAG) affects LLM performance and prevents hallucinations. Empirical results show that our method is more robust and produces fewer hallucinations than vanilla RAG.

• Dependence on the Quality of Dense Document Embeddings

It is widely acknowledged that the performance of the modality encoder heavily impacts downstream performance, and xRAG is no exception. However, xRAG is not limited to only strong embedding models. In Section 6.2, we demonstrate the effectiveness of xRAG across various embedding models, showing that even with a four-year-old DPR model, xRAG still yields over 8% improvement.

• 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?

According to Claude Shannon's information theory, one classic estimate of the amount of information per English word is 11.82 bits per word [1]. This means, in theory, we have enough raw state to encode more tokens with a single embedding without loss of information. In that sense, we can be more efficient when dealing with real-world multi-round, multi-document RAG settings, where prompt length could be up to thousands of tokens. More importantly, in Section 5.2, we have demonstrated that the current xRAG could bring considerable efficiency improvement even on a high-end server GPU.

As for balancing compression and performance, we discuss in Appendix G that xRAG can be combined with a multi-vector or Matryoshka embedding system where the compression ratio is configurable. We believe this is an exciting direction for xRAG.

• 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?

They refer to the same module that connects two modalities.

• The paper mentions using a two-layer MLP as the projector but doesn't explore how different projector architectures might impact performance.

As discussed in Section 3.4, the key design principle is simplicity. Therefore, we chose the simplest and most commonly used projector: a two-layer MLPs. Exploration of more advanced architectures is left for future work.

• 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.

One reason for xRAG's robustness is that it is trained with retrieved context in the second training phase, which improves the LLM's resilience to noisy retrievals. This phenomenon is also observed in [2][3]. Additionally, xRAG can avoid word-by-word repetition seen in RAG when dealing with noisy retrieval results.

• How will the hyper parameter alpha influence the final results?

We tested several values for the hyperparameter alpha ({0.1, 0.5, 1.0, 2.0, 3.0}) and found that 2.0 gave the best results on the validation set. No significant differences were observed when alpha ≥1.

• How to ensure the compressing process filters the irrelevant information rather than the crucial ones?

Our paper, like previous context compression papers, does not claim to achieve lossless compression. It is a balance between efficiency and efficacy. We do not filter out any information; instead, we aim to recover the original information with a single embedding vector. Our novel training paradigm, from paraphrase pretraining to context-aware instruction tuning, is designed to achieve this goal.

• Will sensitive information like numbers or times be kept by the projector after compressing?

xRAG has the ability to retain sensitive information such as numbers and times. As shown in Figure 9 of our paper, xRAG successfully identifies important numbers and answers questions correctly.

Thank you again for your review. We hope that our response alleviates some of your concerns and that you might consider raising the score of our paper.

Best regards,

Authors

[1] Claude Elwood Shannon. “Prediction and Entropy of Printed English.” (1951).

[2] Lin, Xi Victoria, Xilun Chen, Mingda Chen, Weijia Shi, Maria Lomeli, Rich James, Pedro Rodriguez, Jacob Kahn, Gergely Szilvasy, Mike Lewis, Luke S. Zettlemoyer and Scott Yih. “RA-DIT: Retrieval-Augmented Dual Instruction Tuning.” ArXiv abs/2310.01352 (2023): n. pag.

[3] Luo, Hongyin, Yung-Sung Chuang, Yuan Gong, Tianhua Zhang, Yoon Kim, Xixin Wu, Danny Fox, Helen M. Meng and James R. Glass. “SAIL: Search-Augmented Instruction Learning.” Conference on Empirical Methods in Natural Language Processing (2023).

Dear Reviewer,

We greatly appreciate your time and effort in reviewing our paper. Here, we would like to address your concerns point by point:

• Out-of-Domain Generalization

Thank you for bringing up this issue. If we understand your concerns correctly, they can be divided into two parts:

1. The Out-of-Domain Generalization of Dense Retrieval Models: We acknowledge that dense retrieval models face challenges with cross-domain generalization, yet they remain the de facto approach for modern RAG systems [1][2][3]. Importantly, dense retrieval models (such as DRAGON[4] and ColBERT[5]) have demonstrated stronger generalization compared to their sparse counterparts (such as BM25), as evidenced by benchmarks like BEIR[6] and MTEB[7]. These models are all compatible with xRAG.

2. The Out-of-Domain Generalization of xRAG: The ultimate goal of xRAG is to create compatible representations between dense embeddings and their textual form within the LLM representation space using a modality fusion approach. We believe a well-trained projector can transform any text into the embedding space. To verify this, we tested our xRAG model (paraphrase pre-trained in the Wikipedia) in the biomedical domain. Specifically, we tested on PubMedQA[8], BioASQ-Y/N[9], and ChemProt[10] using PubMed[11] as our retrieval corpus to evaluate the cross-domain generalization ability of xRAG. By comparing the first two rows in the table below, we observed that simply adding one document token, xRAG gives considerable improvement over the Mistral-7b baseline, demonstrating the cross-domain generalization of xRAG. Moreover, xRAG can easily adapt to other domains because the first stage of training pipeline: paraphrase pretraining, is a self-supervised process that does not require handcrafted labels. We sample documents from PubMed and pre-train xRAG on it, which further improve the performance as shown in the last row.

• Details about Training and Selection of Dense Models:

  As stated in Section 4.2 of our paper, we selected SFR[12] as our dense embedding model, which, at the time of writing, held the leading position on the MTEB leaderboard. We did not train this model ourselves, allowing us to reuse the offline constructed embeddings and adding ZERO overhead to the existing RAG system, the core of xRAG as detailed in Section 3.4.

• Long-Context RAG:

  RAG is commonly considered a technique to alleviate the need for LLMs to read long contexts by chunking documents into pieces, which is how modern RAG systems operate. Therefore, technically speaking, long-context and RAG represent two different generation paradigms for LLMs. If your concern regarding xRAG involves multiple document chunks, please refer to our general response.

Thank you again for your review. We hope that our response addresses your concerns and that you might consider raising the score of our paper.

Best regards,

Authors

References:

[1] Lin, Xi Victoria, Xilun Chen, Mingda Chen, Weijia Shi, Maria Lomeli, Rich James, Pedro Rodriguez, Jacob Kahn, Gergely Szilvasy, Mike Lewis, Luke S. Zettlemoyer and Scott Yih. “RA-DIT: Retrieval-Augmented Dual Instruction Tuning.” ArXiv abs/2310.01352 (2023): n. pag.

[2] Yu, Yue, Wei Ping, Zihan Liu, Boxin Wang, Jiaxuan You, Chao Zhang, Mohammad Shoeybi and Bryan Catanzaro. “RankRAG: Unifying Context Ranking with Retrieval-Augmented Generation in LLMs.” (2024).

[3] Shao, Rulin, Jacqueline He, Akari Asai, Weijia Shi, Tim Dettmers, Sewon Min, Luke S. Zettlemoyer and Pang Wei Koh. “Scaling Retrieval-Based Language Models with a Trillion-Token Datastore.” (2024).

[4] Lin, Sheng-Chieh, Akari Asai, Minghan Li, Barlas Oğuz, Jimmy J. Lin, Yashar Mehdad, Wen-tau Yih and Xilun Chen. “How to Train Your DRAGON: Diverse Augmentation Towards Generalizable Dense Retrieval.” ArXiv abs/2302.07452 (2023): n. pag.

[5] Khattab, O. and Matei A. Zaharia. “ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT.” Proceedings of the 43rd International ACM SIGIR Conference on Research and Development in Information Retrieval (2020): n. pag.

[6] Thakur, Nandan, Nils Reimers, Andreas Ruckl'e, Abhishek Srivastava and Iryna Gurevych. “BEIR: A Heterogenous Benchmark for Zero-shot Evaluation of Information Retrieval Models.” ArXiv abs/2104.08663 (2021): n. pag.

[7] Muennighoff, Niklas, Nouamane Tazi, Loic Magne and Nils Reimers. “MTEB: Massive Text Embedding Benchmark.” Conference of the European Chapter of the Association for Computational Linguistics (2022).

[8] Jin, Qiao, Bhuwan Dhingra, Zhengping Liu, William W. Cohen and Xinghua Lu. “PubMedQA: A Dataset for Biomedical Research Question Answering.” Conference on Empirical Methods in Natural Language Processing (2019).

[9] Yang, Zi, Yue Zhou and Eric Nyberg. “Learning to Answer Biomedical Questions: OAQA at BioASQ 4B.” (2016).

[10] Peng, Yifan, Shankai Yan and Zhiyong Lu. “Transfer Learning in Biomedical Natural Language Processing: An Evaluation of BERT and ELMo on Ten Benchmarking Datasets.” BioNLP@ACL (2019).

[11] https://huggingface.co/datasets/ncbi/pubmed

[12] https://blog.salesforceairesearch.com/sfr-embedded-mistral/

Dear Reviewer,

We greatly appreciate your time and effort in reviewing our paper. Here, we would like to address your concerns point by point:

• Question about why xRAG could compress a document chunk into a single token

Thank you for this insightful question. We believe this is the core of xRAG, and we want to explain it from two perspectives. First, it is surprising how much information a single embedding can contain. As demonstrated in [1], a single vector of shape 768 can recover 32 tokens with over a 90% exact match. In our experiments, we use an embedding of shape 4096 to cover a document chunk with an average of 180 tokens. Moreover, according to Claude Shannon's information theory, one classic estimate of the amount of information per English word is 11.82 bits per word [2]. This suggests that, in theory, we still have enough raw state to encode more tokens with a single embedding. Second, the overall performance is a joint effect from both Boost Rate and Resilience Rate (as defined in Section 6.1 of our paper) and we have to acknowledge that xRAG does not yet perform comparably with RAG in terms of Boost Rate. Improving the Boost Rate of xRAG while maintaining a high level of Resilience Rate—balancing external and internal knowledge—is a primary focus of our future work.

• About retrieving multiple documents

Please refer to our general response.

Thank you again for your review. We hope that our response alleviates some of your concerns and that you might consider raising the score of our paper.

Best regards,

Authors

[1] Morris, John X., Volodymyr Kuleshov, Vitaly Shmatikov and Alexander M. Rush. “Text Embeddings Reveal (Almost) As Much As Text.” Conference on Empirical Methods in Natural Language Processing (2023).

[2] Claude Elwood Shannon. “Prediction and Entropy of Printed English.” (1951).