Accelerating Inference of Retrieval-Augmented Generation via Sparse Context Selection

First of all, we would like to apologize that we did not present our work clearly enough in our paper. For the questions below, we will update our manuscript to make sure things are clear.

Primary contribution lies in scoring documents and using an existing framework with modified position encodings to achieve parallel encoding

We think this understanding is not accurate. Our novelty not only lies in combining "parallel encoding" with "scoring", but also in demonstrating that we can perform the two separate tasks in just one step with one model without any additional cost. This is significantly different from previous works which either only focus on parallel contexts or use external models to filter the context. To the best of our knowledge, this is novel and no other works have explored a similar approach. Our significant gains in both quality and efficiency in Table 3 should help justify our novelty.

The experiments lack ablation studies for the module, especially an experiment where only Parallel Contexts is added to RAG (for only generation task)

We have this setup as a baseline, which is PCW-RAG.

We have a number of ablations included in the paper, and we are adding an ablation in this rebuttal confirming that Sparse RAG without inter-document attention scores performs almost as well as Sparse RAG with full attention during the generation step, and both perform better than RAG.

Some parts of this work are not clearly written. lines 201 to 204, with line 204 being especially unclear.

We apologize for the confusion. We will update our manuscript to be more clear.

The LLM input format is <Question> <Document> <Control_Assessment>, and output format is <Relevant>, where <Control_Assessment> is a special control token used to instruct the LLM to generate a relevance score, and <Relevant> is a token either “yes” (meaning it is relevant) or “no” (meaning it is not relevant). During inference time, we add “yes” as the <Relevant> token so that the LLM generates the probability that the document is relevant to the current context. This probability is used as the relevance score to filter the documents.

Please let us know if this is still not clear. We are happy to further clarify.

What is the detailed difference between Parallel Context Windows (PCW) and the Parallel Contexts the authors proposed?

There is no difference. We adopted this idea in our approach. But we added the Per Context Assessment and context filtering using the same model to address PCW's limitations that 1/ it creates a quality regression and 2/ it only speeds up the prefill stage, not decoding. Our method, on the other hand, addresses both problems.

what is the {Relevant} token

It is “yes” and “no” which defines relevant in a binary fashion. We will update this in the manuscript.

do ci and di represent the same thing?

No. Di is the retrieved documents. Ci is the context KV cache of the documents.

Is Parallel Contexts used for both assessment and generation?

Yes it is. When we say “Parallel Contexts is used only for the generation task.", we refer to the assessment training task, which only has 1 context, but in both training and inference, in both prefill and generation, we apply a parallel contexts mask. We will clarify this in our new manuscript.

experiment to evaluate Corrective RAG using LLM Labels?

Thanks for this great suggestion. We added this ablation in the table below. The results show that if we adopt our LLM labels with CRAG we see marginal improvements in performance -- but the results still lag behind our approach. Sparse RAG still outperforms CRAG on both quality and efficiency.

This design relies on pretty heavy training data. How does your method work if such training data is not available, or when your data and applications do not match the statistics in your training data?

In this paper, we proposed a solution to this challenge. Specifically, we leveraged LLM critiques to generate synthetic data when suitable training data was unavailable. As demonstrated in our evaluation, this approach proved to be highly effective, outperforming silver-label methods in the case of PopQA. While testing in out-of-distribution scenarios is not the primary focus of this study, we emphasize the potential of employing LLM critiques more effectively, given their demonstrated ability to assess document relevance with notable accuracy [1]. Additionally, by collecting as much available data as possible and employing ensemble techniques, we can derive a distribution that is more generalizable across domains and datasets lacking sufficient training data.

[1] Retrieval Augmented Generation or Long-Context LLMs? A Comprehensive Study and Hybrid Approach. EMNLP 2024.

semantic analysis on query intent and nor does any concrete semantic analysis on the text to be retrieved.

Our design intentionally avoids performing explicit semantic analysis of query intent and the text to be retrieved to prioritize simplicity and computational efficiency. We acknowledge that incorporating semantic analysis could potentially enhance the retrieval accuracy, especially in cases where query intents or ambiguous language need to be interpreted. Such analysis might better capture subtle relationships between queries and documents, leading to more precise matching. However, there are also potential drawbacks to integrating semantic analysis. It could significantly increase the system’s complexity and computational overhead, potentially impacting scalability.

In this study, we chose not to include explicit semantic analysis in our design because our primary focus was on evaluating the effectiveness of leveraging LLM critiques in achieving a balance between the efficiency of RAG inference. We aimed to establish a baseline performance that could later be extended or refined in future iterations. Nevertheless, we agree that discussing this tradeoff would provide valuable context, and we will include a discussion on the potential impact of semantic analysis and the rationale behind our design choices.

elaborate on how the Sparse RAG approach specifically handles such time-sensitive or ambiguous queries.

Thanks for this question. Section B in the Appendix addresses the absence of sufficient context to correctly answer the question in the widely used public datasets we experimented on. For consistency with the literature, we did not modify the queries or the retrieved contexts, but we do expect that the ambiguity is impacting our results. In many cases, retrieved contexts would be assessed differently for relevance depending on the true intended meaning of the question, and different answers would be expected. It would be an interesting future exploration to augment the datasets to resolve such ambiguity, as it would likely improve the accuracy of our relevance labels and improve performance overall.

For missing time or location context, you could augment the datasets by simply adding new contexts with the missing information or appending the context to the query. For example, to handle the question "Who is the current president of the United States?", you could add a context "Today is November 18, 2024" or you could modify the query to "Who is the current president of the United States on November 18, 2024"?. In the paper, we showed that Sparse RAG performs comparably to RAG on the HotpotQA multi-hop reasoning dataset, so we expect that Sparse RAG could effectively leverage information from new contexts in the generation step, but appending such temporally specific context to the query itself would likely yield the best quality, as all contexts attend to the query.

In real world cases where the question is ambiguous, the most desirable behavior would be to request clarification before further action. In an experimental environment, we could use LLMs to rewrite the queries to contain disambiguating details without giving away the answers. Alternatively, one could simulate requesting clarification by adding an additional binary ambiguous/clear classification task for each question and augmenting the dataset with disambiguating information similar to the approach described above. Then, in evaluation, the model would classify the question's ambiguity, then request more information if necessary, then perform per-context assessment, and then perform generation.

We will add some discussion of this to the new manuscript Appendix.

Dear Reviewer 4HHs,

As the rebuttal deadline approaches, please kindly check the papers' discussion threads and respond to the authors' rebuttals. If you haven't had a chance to respond yet, I’d greatly appreciate your input soon. Your insights are invaluable to the authors and the review process.

Thank you for your effort and support!

Best regards,

Area chair

The paper does not provide experimental comparisons to validate that the absence of inter-document attention scores during inference does not lead to a significant decline in the quality of responses.

We would like to clarify that our experiments already show that the absence of inter-document attention scores during inference does lead to a quality regression. PCW-RAG, which is one of our baselines, uses inter-document attention masking without our document assessment training and filtering. In Table 3 in the paper, PCW-RAG has a clear drop in quality, especially for HotpotQA which requires multi-hop reasoning. Sparse RAG, on the other hand, which also omits inter-document attention during inference, remedies this quality drop by removing irrelevant contexts from the generation step, achieving better quality than RAG.

Moreover, as part of this rebuttal, we added an additional ablation confirming that the Sparse RAG without inter-document attention scores performs almost as well as Sparse RAG with full attention during the generation step, and both perform better than RAG.

Additional complexity and require careful tuning to ensure effectiveness.

Although our approach does add some additional complexity and does require further training, the changes required are minimal and only affect the position encoding and the attention mask. This is little work to implement and adds little additional complexity in terms of training efficiency. More importantly, our approach minimizes computation and requires fewer resources during inference time, which could be seen as a complexity decrease.

Regarding “careful tuning”, our experiments show that a natural combination of the two training tasks–including all retrieved contexts as examples for the classification task and all queries for the generation task–works well. Adjusting the mixture weight does not cause obvious differences in performance. This suggests that tuning is pretty stable.

Does reusing the KV Cache during batch inference cause a decline in the quality of responses when the document-level attention scores are missing?

Yes, it does, but as we explained in our response to your first point above, the document assessment step improves quality beyond that of full-attention RAG with large efficiency gains.

Why does the use of an external classifier in the CRAG method increase so much latency?

The T5 model compared here is a T5 large, which is a 750M size model. On the contrary, the LLM profiled here is the XXS size of Gemini (scale of 1000M), which can run on-device. Thus, this classifier is not small in our settings. Also, note that our baseline model Gemini has a different implementation from the old-style LLM (T5). For example, T5 has a much larger number of attention layers than the XXS version of Gemini. Gemini XXS is also implemented with flash attention and other optimizations which improve speed, as opposed to T5’s model, which does not include these performance optimizations.

Is the training data for the classifier consistent with the PCA task

Yes, the classifier is also trained with the PopQA dataset so that the methods can be fairly compared.

Dear Reviewer 5W8k,

As the rebuttal deadline approaches, please kindly check the papers' discussion threads and respond to the authors' rebuttals. If you haven't had a chance to respond yet, I’d greatly appreciate your input soon. Your insights are invaluable to the authors and the review process.

Thank you for your effort and support!

Best regards,

Area chair

ablation study to evaluate the impact of concatenating selected KV caches instead of the corresponding raw contexts

Thanks for this great suggestion. We share results on the PopQA dataset in the table below, where we used a full-attention model to process the same selected context. Using full attention provides a slight quality improvement, but Sparse RAG quality is very close and more efficient.

Comparison with KV Cache Compression Methods

Thank you for the thoughtful suggestion. We acknowledge the potential insights that such a comparison could provide. However, our primary focus in this work is to demonstrate the effectiveness of SparseRAG’s context selection strategy in enhancing efficiency during prefill and decoding stages, which we believe represents a substantial and novel contribution in itself. While KV cache compression methods typically operate under different paradigms and are training-free, integrating a detailed comparison would require addressing additional factors such as varying use cases and performance metrics. Conducting such an analysis is a non-trivial extension that we believe is better suited as a direction for future research.

We appreciate your suggestion, and it will certainly inspire further exploration of these ideas in subsequent work. For now, we hope the presented results of comparing LLMLingua and PCW-RAG sufficiently illustrate the utility and advantages of SparseRAG in its intended scope.

Some parts in the presentation are not clear

We apologize for any unclear presentation. We will provide clarification in the new manuscript. To address your questions, 1/ The rating tokens are “yes” and “no”. 2/ Yes, “assess” and “generate” are special tokens. 3/ The position encoding is "parallel incremental”. For example, in “0 1 2 3 4 5 3 4 5 6 10 11 12”, (0 1 2) is the question, (3 4 5) is the first context, (3 4 5 6) is the second context, and (10 11 12) is the response. 4/ During inference, the assessment does use the cached representations in input.

Context reduction methods have been widely studied for RAG

Thanks for this suggestion. We understand the importance of comprehensive comparisons to situate our work within the broader landscape of research. We would like to highlight that our baselines already include two effective and popular approaches specifically focused on context reduction, i.e., CRAG and LLMLingua. These baselines were selected due to their relevance and strong performance in addressing the challenges of context compression. By including these methods, we ensure a robust and meaningful evaluation of our proposed approach. While we agree that more baselines could theoretically provide further insights, we believe that our current selection is sufficient to demonstrate the effectiveness of our method, particularly given the computational and practical constraints of extensive comparative studies. We remain open to expanding this comparison in future work as additional SOTA methods emerge or if a specific context warrants further exploration.

Document assessment method needs to be validated separately

We have validated our document assessment method separately in several ways. 1/ We created a golden dataset with 3 human annotators on a subsample of the Natural Questions dataset and compared F1-score for different LLM labeling methods to select the best one. Details are in Section 4 and the Appendix. 2/ We performed an end-to-end evaluation of our label quality compared to silver labels from the CRAG classifier training data (derived heuristically). We showed that all quality and efficiency metrics improved by using our Sparse RAG LLM-generated labels.

We also add an ablation in this rebuttal using our document assessment and filtering method followed by generation with full attention. It shows the isolated effects of our document assessment method. The method provides quality improvements over standard RAG with and without cross-document attention. We are open to performing further assessment if you have specific requests.

SparseRAG needs to be further validated under other advanced CoT, long chain of reasoning, and instruction following tasks.

Thanks for this suggestion. Our proposed method can be considered an instantiation of CoT. Specifically, Sparse RAG first determines the relevance of each document to the query. Then, it performs filtering based on the document assessment result. Finally, it utilizes the document caches for generation. This CoT-inspired process leverages the reasoning capabilities of LLMs. Experimentally, we evaluated Sparse RAG on the HotpotQA dataset, which requires multi-hop CoT reasoning. The results demonstrate that Sparse RAG achieves comparable quality to RAG while offering significantly improved efficiency.

Document assessment can also be parallelizable, using restricted cross-attention methods?

We would like to clarify that our document assessment is not performed one-by-one; it is already parallelizable using restricted cross-attention because documents are treated independently by masking inter-document attention. This improves efficiency for both server-side and on-device inference.

Dear Reviewer p444,

As the rebuttal deadline approaches, please kindly check the papers' discussion threads and respond to the authors' rebuttals. If you haven't had a chance to respond yet, I’d greatly appreciate your input soon. Your insights are invaluable to the authors and the review process.

Thank you for your effort and support!

Best regards,

Area chair