Accelerating Inference of Retrieval-Augmented Generation via Sparse Context Selection

## Abstract

Abstract Large language models (LLMs) augmented with retrieval exhibit robust performance and extensive versatility by incorporating external contexts.
However, the input length grows linearly in the number of retrieved documents, causing a dramatic increase in latency.
In this paper, we propose a novel paradigm named Sparse RAG, which seeks to cut computation costs through sparsity.
Specifically, Sparse RAG encodes retrieved documents in parallel, which eliminates latency introduced by long-range attention of retrieved documents.
Then, LLMs selectively decode the output by only attending to highly relevant caches auto-regressively, which are chosen via prompting LLMs with special control tokens.
It is notable that Sparse RAG combines the assessment of each individual document and the generation of the response into a single process.
The designed sparse mechanism in a RAG system can facilitate the reduction of the number of documents loaded during decoding for accelerating the inference of the RAG system.
Additionally, filtering out undesirable contexts enhances the model’s focus on relevant context, inherently improving its generation quality.
Evaluation results of two datasets show that Sparse RAG can strike an optimal balance between generation quality and computational efficiency, demonstrating its generalizability across both short- and long-form generation tasks.

## 1 Introduction

Large language models (LLMs) have attracted increasing attention and exhibited impressive abilities to understand instructions and generate fluent outputs in natural language [, , , ].
Nevertheless, LLMs inevitably manifest hallucinations [] due to their struggle with factual errors and inability to secure the accuracy of generated text solely by the parametric knowledge they encapsulate [, ].
Feeding the source of truth to LLMs in the format of retrieved context segments [] alleviates this problem. The technique is widely known as Retrieval-Augmented Generation (RAG) [, , ].

Although the RAG framework is empirically shown to be effective, it can be expensive to scale up. This is because it requires prepending relevant documents retrieved from an external knowledge corpus to the queries []. As a result, the input length grows linearly in the number of documents, causing a dramatic increase in latency when using a standard Transformer whose the latency scales quadratically with the input length.
Some prior works such as Fusion-in-Decoder (FiD) [] and Parallel Context Windows (PCW) [] have proposed to alleviate this issue. Yet these methods fail to strike an optimal balance between generation quality and computational efficiency.
FiD was originally designed for the encoder-decoder architecture, and thus is not compatible with currently prevalent decoder-only architectures without significant changes. While
PCW can be applied to decoder-only LLMs, it only speeds up the model pre-filling and still incurs high latency since the whole context window cache is still being attended to when decoding each token.
Moreover, the heavy reliance of generation on the retrieved knowledge raises significant concerns about the model’s behavior and performance in scenarios where retrieval may fail or return inaccurate results [].
A typical approach for mitigating this issue is to rely on an external classifier to rank or filter the documents before prepending them to the input [], but this process requires extra model calls which adds new complexity to inference.

Figure: Figure 1: An overview of Sparse RAG at inference. Each of the retrieved documents will be assessed by LLMs to decide whether to keep or drop by estimating a relevance score. Then, load the documents that are considered to keep for generation.
Refer to caption: x1.png

In light of the issues above, we propose a novel paradigm named Sparse RAG.
It generally operationalizes by massive pre-filling where the key-value cache was generated through a single forward pass of the input tokens and selective decoding where the output is generated by attending to highly relevant tokens auto-regressively.
Previous works where the length of the retrieved contexts during pre-filling are equal to that during decoding are called *dense*-RAG in this paper.
Different from that, Sparse RAG allows the pre-filled context to be significantly larger than the decoding context, where partial low-quality context is dynamically dropped according to their relevance to the input query. Furthermore, Sparse RAG combines the assessment of each individual context and the generation of the response into a single process.
Special control tokens are used to prompt LLMs to assess each retrieved context for picking high-quality ones.
Then only the key-value caches of those picked contexts are loaded into LLMs for decoding with another control tokens.

The design of such a sparse mechanism in RAG system brings two unique advantages.
Firstly, by reducing the number of key-value cache loads during the decoding process, LLMs can achieve lower latency as they are typically constrained by memory usage.
Besides, filtering out undesirable contexts enhances the model’s focus on relevant context, inherently improving its generation quality.
To demonstrate the effectiveness and efficiency of the proposed method, we evaluate on two datasets of PopQA [], and Biography [].
Experimental results show that Sparse RAG can achieve the similar or better performance while keep much better latency compared with standard dense-RAG or PCW-RAG, demonstrating its generalizability across both short- and long-form generation tasks.

## 2 Related Work

##### Retrieval-Augmented Generation

RAG is a family of techniques for generating output structures while using retrieved nearest-neighbor structures as a reference. It typically involves two stages: retrieval and generation.
Retrieval finds most similar contexts based on BM25 or learned embedding, where the context could be represented as token embedding [, ], dense embedding [] or raw text [, , ].
Once those contexts are retrieved, different architectures are leveraged to incorporate them into the model.
Popular approaches includes concatenation [, ] and cross-attention [, ].
In recent years, the LLM architecture evolves towards a single decoder only model with significant larger sizes. To this end, concatenation of raw text [] is becoming popular for being simple and practical.
Many advanced approaches have been developed on top of it in recent years.
Self-RAG [] is proposed to selectively retrieve knowledge and introduce a critic model to decide whether to retrieve.
Yoran et al. [] designed an NLI model to identify the irrelevant context and improve robustness.
SAIL [] is tuned on instructions to insert retrieved documents before instructions.
Toolformer [] is pre-trained for calling APIs such as Wikipedia.
Jiang et al. [] actively anticipate future content and decide when and what to retrieve in long-form generation.
CRAG [] makes the attempt to explore and design corrective strategies for RAG to improve its robustness of generation. Specifically, an external T5 model is trained and used to determine the usefulness of the retrieved context.
Generally, these approaches target on exploiting retrieval as a useful tool to augment generation or whether retrieval is necessary.

##### Efficiency in RAG

The efficiency of LLM inference is a widely explored research area, where different categories of approaches have been studied. Some works are towards architecture level acceleration such as efficient attention [], Mixture of Expert [], Transformer-alternative architectures [], etc. While some other works are algorithm-level acceleration like quantization [] or speculative decoding [].
Nevertheless, most previous works are not RAG specific.
Recently, RAG Cache [] was proposed as a multilevel dynamic caching system tailored for RAG from the system perspective. Another stream of the work try to process each document individually so that no cross document attention computation is required, such as FiD [] or PCW [].
FiD encoded each retrieved passage independently from other passages and decoded by attending over the concatenation of the resulting representations of all the retrieved passages.
PCW carved a long context into chunks (“windows”), restrict the attention mechanism to apply only within each window, and re-use the positional embeddings across the windows.

Comparison with previous works that are the most relevant to our work is illustrated in Table 1.
This work aims to strike an optimal balance between generation quality and computational efficiency.
It is notable that the extra classifier in CRAG requires maintaining extra model with more complex serving infrastructure. When there is $N$ contexts retrieved. There is totally $N+1$ model runs.
Our work also relies an classification process to refine the retrieved documents. However, our approach relies on the “internal” classification process that is aligned with the generation process. The total model run in our case is $1$.

**Table 1: Comparisons with existing RAG-related works.**
| Approach | Corrective | No extra model | Prefill efficiency | Decode efficiency |
| --- | --- | --- | --- | --- |
| RAG [] | No | Yes | No | No |
| Corrective RAG [] | Yes | No | No | No |
| PCW RAG [] | No | Yes | Yes | No |
| Sparse RAG | Yes | Yes | Yes | Yes |

## 3 Sparse RAG

Sparse RAG is designed for the decoder-only model architecture, which is the default case for most popular LLMs.
Figure 1 presents an overview of Sparse RAG at inference, which designs the document assessment strategies to improve the robustness of generation.
The key hypothesis of our approach is that the RAG task and per context assessment is a similar tasks and the model could do both tasks in one shot.
We found simple and effective approaches to enable them.

### 3.1 Training Process

Our work is based on the assumption that a certain amount of RAG training data is accessible, which allows us to effectively tailor and adapt existing LLMs to our specific needs.
In the training phase, we integrate an additional Per Context Assessment (PCA) task into the training mixture.
This PCA task is included alongside the primary RAG task to enhance the model’s ability to assess and respond accurately within various contexts.
By incorporating the PCA task, we aim to improve the overall performance and contextual understanding of the LLMs, ensuring they are more adept at handling a diverse range of scenarios presented during their use.

##### Data Augmentation with LLMs

For typical RAG data, one question-answer pair can be mapped to multiple retrieved contexts using either BM25 or existing standalone retriever. However, there are cases where no golden label regarding the quality of every retrieved context is available.

To collect these missing labels, we leveraged two off-the-shelf LLMs [, ]–PALM2 and Gemini–to get the assessment for each context. We observe empirically that a second round of prompting for critique, especially using a different model from the initial round, ensures the best quality labels. We provide our prompts in Table 9 in the Appendix. We compare different model combinations for labeling to human ground truth labels in Section 4.

##### Multitasking Data Format

The LLM is trained on a mixture of two types of tasks, “Rate” assessment and answer generation. Specifically, we format the inputs and outputs of the two task types as

- •
Assessment: {Question}{Context}{Control_Assessment}{Rate}
- •
Generation: {Question}{Context_1}…{Context_N}{Control_Generation}{Answer}

where Rate and Answer are targets for the generative tasks while everything before them are inputs. Control_Assessment and Control_Generation are fixed tokens to ensure the LLM can differentiate the two tasks.

##### Parallel Contexts

Since each context is rated independently in PCA task, we introduce this independence in the primary RAG task as well so two tasks could reuse KV cache during inference time.
Thus for the generation task, we enforce no cross attention between different two contexts following the Parallel Context Windows [] to process the input for training.
Specifically, we modify two things in the standard LM training process. First, we
change the attention masks to be block-wise, and restrict that context_i and context_j will not attend to each other. Question, Context_i and Control_Generation still follow the default attention mechanism, i.e. in the causal attention case, the latter could attend to the previous ones.
Second, we adjust the position encoding to “mimic” the situation that Context_i is not visible to context_j that each context will have its own incremental positional ids. However, the position id of Control_Generation is designed to back to normal, where its position considers all the length of contexts.

### 3.2 Inference Process

Given the question and retrieved contexts, Sparse RAG handles the assessment task and generation task in one single pass.

##### Per Context Assessment

Similar to the training process, when pre-filling the KV cache, each retrieved context is treated independently. Then these KV cache will be used to score the input, which is the concatenation of the token of Control_Assessment and the token of “Good”. The score is the probability of this tokens as the next tokens.

##### Generation

The generation uses a filtered KV cache, where only $K$ out of $N$ KV cache will finally be loaded. We use a simple threshold based filtering approach: we drop the context when its score is less than $sigma$. Once those caches are loaded, Control_Generation will be used to prompt the model to generate the answers.

## 4 Evaluation of Per Context Assessment

### 4.1 New Benchmark: Natural Question Per Context Assessment

We isolated a subset of 50 questions, each with 10 retrieved contexts, from the Natural Questions dataset. We assigned 3 raters to each question-context pair from a pool of 7 raters and provided the instructions in Section A.2.

We aggregated responses for all 3 raters for each context, selecting the majority decision 0 or 1 for each context. We found that raters unanimously agreed on 351 out of 500 context, with  30% of the documents considered relevant. For questions where raters were not unanimously decided, a specialist rater was assigned to investigate more carefully and set the best label to correct mistakes of the other raters.
This resulted in 6 additional documents considered relevant out of the entire dataset, boosting the portion of relevant documents to 31% and slightly increasing alignment with the auto-rater approaches (average F-score increase of 1.4% across auto-rater methods using these corrections as the ground truth).

### 4.2 LLM rater comparisons

We tested several different LLM-based automatic labeling methods–different combinations of models and prompts–for creating training data for the classifier in Sparse RAG. We compared several of these auto-rater approaches by creating a ground-truth relevance dataset using human labeling.
The auto-rater comparison using the revised human labels as the ground-truth is shown in Table 2. We find that combining two different models in two rounds–initial prompting and critique–provides the labels that are most closely aligned with the human labels. We hypothesize that the different representations learned by two different models are able to capture the most nuance in the input sequences, leading to better relevance judgements. We also observe that Gemini Ultra appears slightly less effective at critiquing model outputs than PALM2 XL.

**Table 2: Auto-rater comparison to ground truth.**
| Auto-labeling method | Average F1 | F1 Label 0 | F1 Label 1 |  |
| --- | --- | --- | --- | --- |
| Rater model | Critic model |  |  |  |
| PALM2 XL | n/a | 0.729 | 0.765 | 0.694 |
| PALM2 XL | PALM2 XL | 0.781 | 0.820 | 0.741 |
| Gemini Ultra | n/a | 0.761 | 0.807 | 0.716 |
| Gemini Ultra | Gemini Ultra | 0.704 | 0.747 | 0.660 |
| PALM2 XL | Gemini Ultra | 0.728 | 0.776 | 0.680 |
| Gemini Ultra | PALM2 XL | 0.821 | 0.861 | 0.781 |

## 5 Evaluation of Sparse RAG

### 5.1 Benchmarks and Metrics

##### PopQA

is a large-scale open-domain question answering (QA) dataset, consisting of 14k entity-centric QA pairs. Each question is created by converting a knowledge tuple retrieved from Wikidata using a template. We follows the setup from  [] and use Contriver [] to retrieve the related contexts.
Since PopQA is lacking the per-context assessment label, we adopted the “Gemini + PALM2” combination to get training labels. We split the dataset into training, validation and test set with 8:1:1 ratio.
Since the answer is usually short, we report Exact Match (EM) and F1 scores.

##### QMSum

[] is a human-annotated benchmark for query-based multi-domain meeting summarization task, which consists of 1,808 query-summary pairs over 232 meetings in multiple domains. To adopt it in the RAG domain, we divide each conversation into different contexts where the average context contains 300 words, and filter the data point where the conversation containing ground truth label is split into multiple contexts.
Note that this dataset has human labeled per-context assessment that we could leverage during training.
Eventually we get $250$ training data, $70$ validation and $77$ test set. The target of this data is longer and we report RougeLSum and F1 scores.

### 5.2 Baselines

##### RAG

We evaluated the standard RAG [] where an LM generates output given the query prepended with all the top ranked documents using the same retriever as
in our system.

##### PCW RAG

We applied the Parallel Context Windows [] to the RAG process, where no cross attention is applied between documents.

##### Corrective RAG

It is the same with RAG except that an external T5-XXL classifier is trained using heuristic labels []. This classifier is used to process all the document and decide the rank. Note to facilitate the fair comparison, we did not adopt the "web search" feature of this paper.

### 5.3 Experimental Configuration

The base LLMs used in the paper Gemini [].
Although our approach could be applied to different training stage of the model, we apply LoRA tuning [] to enforce alignment on top of the foundation LLMs due to its low resource requirement and wide usage. Note that the same LoRA tunning is applied to Sparse RAG and all baselines. In all our experiment, we apply LoRA in self-attention and use the default rank as 4. By default, we use the XXS size which could be running on-device.

During training, we use $64$ Tensor Processing Units (TPU) V3 chips for PopQA while use $128$ Units for .
The batch size is $64$.
We use the Adafactor optimizer [] with a learning rate of $0.003$. The training dropout rate is $0.05$.
We leverage the metrics of validation set to pick the best checkpoint.
During inference, the temperature is set to $0.5$. Unless specifically noted, we use sampling decoding with sample number $1$ for our experiments.

### 5.4 Inference Setup and Metrics

Evaluation of Sparse RAG was conducted on a Samsung S21 Ultra, utilizing the device’s CPU to assess real-world performance on a relatively mid-tier smartphone compared to the latest flagship models. Inference configuration consisted of fixed token lengths for queries, contexts and generated responses. This setup allows for evaluating the system’s efficiency and effectiveness under resource constraints typical of mobile devices, providing insights into its practical applicability for on-device question-answering tasks.

Specifically, the overall inference process considers two stages.

- •
Prefill stage:
For baseline RAG Model, we measure the total time taken to process all input tokens (question and all contexts).
For PCW RAG and Sparse RAG models: We take advantage of these models’ ability to cache question processing. We first measure the time to process the question alone. Then, we measure the time to process different contexts with the pre-processed kv-cached question. We use Encoding Speed (ES) (tokens per second (t/s)) to measure the efficiency of prefill stage.
- •
Decoding stage:
To assess decoding speed comprehensively, we generate output sequences using the same question but varying the amount of relevant context data considered, ranging from the top 1 most relevant document to the top $K$. For each context size, we produce output sequences of different lengths. This systematic approach allows us to evaluate the impact of both context size and response length on decoding speed. We use Decoding Speed (DS) (tokens per second (t/s)) to measure the efficiency of decoding stage.

### 5.5 Main Results

We report the results with both short form answer (PopQA) and long form answer (QMSum) in Table 3, where both quality and latency metrics are shared. “K” is the number of chosen memories.

Notably, our proposed approach achieves the best quality while being the most efficient during inference compared to other “Dense RAG” approaches.
It can be seen that Sparse RAG shares the same pre-filling efficiency with PCW, due to the parallel context encoding. However Sparse RAG can achieve significant better decoding efficiency than standard RAG and PCW RAG.
In detail, out of 20 retrieved contexts, the Sparse RAG has average of $7.84$ context for PopQA and $4.45$ context for QMSum. This leads to almost double or even triple decoding speed.
Meanwhile, Sparse RAG achieves higher quality metrics than the dense counterparts, which shows that Sparse RAG effectively filtered noisy and irrelevant contexts.

On the other hand, it is interesting to compare Corrective-RAG(CRAG) with Sparse RAG on PopQA task since Corrective-RAG trained their own external classifier with T5 XXL based on PopQA data. Results suggested that our approach, albeit using an “in-place” classifier, is outperforming the CRAG with external classifier. Note the encoding and decoding speed of CRAG is not comparable because it involves multiple model running from the classifier.

**Table 3: Quality & Efficiency Tradeoff. For both short form and long form generation tasks, Sparse RAG achieves both higher quality and efficiency compared to “dense” RAG approaches.**
|  | ES | PopQA | QMSum |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | EM | F1 | K | DS | F1 | RougeLSum | K | DS |  |
| RAG [] | 56.28 | 65.43 | 69.99 | 20 | 6.65 | 21.43 | 18.20 | 20 | 6.64 |
| PCW-RAG[] | 147.58 | 65.04 | 69.54 | 20 | 6.65 | 20.18 | 16.95 | 20 | 6.64 |
| CRAG [] | - | 66.52 | 70.99 | 8.90 | - | - | - | - | - |
| Sparse RAG | 147.58 | 67.17 | 71.16 | 7.84 | 12.28 | 23.96 | 20.10 | 4.45 | 16.05 |

### 5.6 Analysis

#### 5.6.1 Impact of Confidence Threshold

To show how the metrics changes with different thresholds, the performance of Sparse RAG corresponding to different thresholds is illustrated in Table 4.
As the threshold gradually increases, the system will filter out more contexts, which would always reduce the latency during inference.
In terms of generation quality, it can be seen that the performance of Sparse RAG was significantly improved as thresholds increased at the beginning, which shows the effectiveness of filtering out irrelevant contexts.
Then, the performance was stable and dropped slightly.
The reason might be that some truly relevant contexts are accidentally filtered out.

**Table 4: Sampling a few confidence threshold on PopQA and QMSum. Higher confidence threshold means filtering more memories. Efficiency wise, filtering more always helps; While in terms of quality.**
|  | PopQA | QMSum |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Threshold | EM | F1 | K | DS | F1 | RougeLSum | K | DS |
| 0.05 | 66.95 | 70.97 | 9.75 | 10.61 | 22.85 | 19.49 | 7.92 | 11.92 |
| 0.1 | 66.84 | 70.66 | 8.72 | 11.70 | 23.78 | 19.98 | 6.68 | 12.89 |
| 0.15 | 67.17 | 71.16 | 7.84 | 12.28 | 23.43 | 19.66 | 5.77 | 13.01 |
| 0.2 | 66.77 | 70.54 | 7.13 | 12.88 | 23.2 | 19.79 | 5.05 | 14.54 |
| 0.25 | 65.75 | 69.64 | 6.56 | 13.00 | 23.96 | 20.1 | 4.45 | 16.05 |
| 0.3 | 63.86 | 68.2 | 5.98 | 13.08 | 23.84 | 19.99 | 3.93 | 16.38 |

#### 5.6.2 Silver Labels vs LLM Labels

In Corrective RAG, the T5 model was trained with silver labels that comes from the title matching []. We collect the same silver labels to replace the LLM labels and train the model with this new dataset. The comparison result is shown in Table 5.

From these results, we can observe that the quality of the labels generated by LLMs is higher than that of the silver labels from Yan et al. [].
We hypothesize that the superior quality of the LLM-generated labels is attributable to our methodology, which involved a two-round process of soliciting responses from two different LLMs.
By engaging two distinct models, we likely enhanced the robustness and accuracy of the labels through a form of cross-validation, thereby mitigating potential biases or errors that might arise from relying on a single LLM.

**Table 5: Comparing silver labels with LLM labels on PopQA.**
| Approach | EM | F1 | K | DS |
| --- | --- | --- | --- | --- |
| Sparse RAG | 67.17 | 71.16 | 7.84 | 12.28 |
| w/ silver label | 66.97 | 71.05 | 8.26 | 11.99 |

#### 5.6.3 Ablation of Prefill Documents

One question one may have is that the better performance may rely on the “massive” prefilling but in practical situation, we can retrieve smaller amount of context which have better quality.
To this end, we compare different numbers of the prefill document: 10 vs 20 for PopQA, and show the results in Table 6.
Even in the setting lower prefilled case, the previous conclusion of Sparse RAG remains. It achieves better quality with higher speed of decoding.
It makes sense that the gap between dense RAG and sparse RAG is relative smaller at 10 prefill documents. This is because the headroom is smaller as the top 10 documents naturally has less to filter and make it less “sparse”. Yet our approach still achieves best trade-off, showing the capability to generalize with different numbers of input documents.

**Table 6: Ablation on Different number of prefill documents for PopQA.**
| Approach | Prefill Documents | EM | F1 | K | ES | DS |
| --- | --- | --- | --- | --- | --- | --- |
| RAG | 10 | 64.66 | 68.67 | 10 | 80.74 | 10.31 |
| PCW RAG | 10 | 63.9 | 68.58 | 10 | 147.48 | 10.31 |
| Sparse RAG | 10 | 65.86 | 70.2 | 7.79 | 147.48 | 12.33 |
| RAG | 20 | 65.43 | 69.99 | 20 | 56.28 | 6.65 |
| PCW RAG | 20 | 65.04 | 69.95 | 20 | 147.58 | 6.65 |
| Sparse RAG | 20 | 67.17 | 71.16 | 7.84 | 147.58 | 12.28 |

#### 5.6.4 Inference Efficiency Ablations

Focusing on the decoding stage, we delve into the performance differences when varying the number of retrieved memories (i.e., top-K documents) and the length of the generated response. This analysis provides valuable insights into the computational overhead associated with incorporating additional context and generating more extensive outputs, ultimately shedding light on the practical trade-offs for real-time question-answering on resource-constrained devices.

As illustrated in Fig 2(b), we observed that the baseline method requires over 50%(varies from 52% to over 55% depending on number of output tokens) more time to generate outputs of varying lengths compared to the Sparse RAG approach. This significant disparity in generation speed is primarily attributed to the increased computational burden incurred by the baseline method when considering a larger number of memories. As shown in Fig 2(a), this heightened demand for computational resources results in a notable slowdown in terms of tokens per second (t/s). This observation underscores the efficiency advantages offered by Sparse RAG, especially in scenarios where a substantial amount of context needs to be considered during the decoding process.

Figure: (a) Decoding speed with different number of memories
Refer to caption: extracted/5619503/figures/graph1-1.png

#### 5.6.5 Ablation on Foundation Model Size

We applied Sparse RAG to different sizes of LLMs, specifically testing it on Gemini XS and Gemini XXS.
The results of these experiments are presented in Table 7.
The findings demonstrate that Sparse RAG is compatible with various foundation models, effectively adapting to different model sizes.
Notably, with a reduced amount of decoding caches, Sparse RAG is capable of achieving the highest quality results.
This indicates that Sparse RAG maintains its efficiency and effectiveness across different foundation models, making it a versatile approach for various LLM configurations.

**Table 7: Ablation on Different LLM model sizes.**
| Approach | Model Size | EM | F1 | K |
| --- | --- | --- | --- | --- |
| RAG | XS | 66.52 | 70.87 | 20 |
| PCW RAG | XS | 65.75 | 70.37 | 20 |
| Sparse RAG | XS | 68.26 | 72.26 | 6.27 |
| RAG | XXS | 65.43 | 69.99 | 20 |
| PCW RAG | XXS | 65.04 | 69.95 | 20 |
| Sparse RAG | XXS | 67.17 | 71.16 | 7.84 |

#### 5.6.6 Using Golden Context Label during inference

Since QMSum provides golden per-context labels, we leverage these labels during inference to evaluate the upper bound performance under the condition of perfect per-context assessment.
By utilizing the highest quality labels available, we aim to determine the best possible outcomes that our model can achieve.
The results of this experiment, which highlight the performance ceiling under ideal labeling conditions, are presented in Table 8.
This approach allows us to understand the maximum potential of our model when supplied with optimal input data, thereby offering insights into its ultimate capabilities.

**Table 8: Trying golden labels on QMSum.**
| Approach | F1 | RougeLSum | K | DS |
| --- | --- | --- | --- | --- |
| Sparse RAG | 23.96 | 20.1 | 4.45 | 16.05 |
| + golden label | 26.76 | 21.93 | 1.13 | 21.16 |

## 6 Conclusion

This paper presents Sparse RAG to address the challenges of increased input length and latency.
Through a novel approach of massive pre-filling and selective decoding, Sparse-RAG efficiently manages the key-value cache of retrieved documents, allowing the LLMs to focus on highly relevant tokens.
This selective attention mechanism not only reduces the computational burden during inference but also enhances the generation quality by filtering out irrelevant contexts.
The evaluation on two datasets validates Sparse RAG’s ability to achieve a balanced trade-off between high-quality generation and computational efficiency, proving its versatility and effectiveness for both short- and long-form content generation tasks.
This innovative paradigm showcases the potential for improving LLM performance in various applications by optimizing context management and inference processes.

## Limitation and Future Work

Sparse RAG requires additional tuning on top of the LLMs, which is not a plug-and-use approach.
Under cases where per-context assessment label is missing, the prompts used for LLM rating may need to be adjusted for different use cases.
These adjustments ensure that the model can still function effectively despite the lack of specific context labels.
Furthermore, future research will explore Sparse-RAG in multimodal contexts, investigating how Sparse-RAG can handle and integrate information from multiple types of data to improve its performance and applicability across diverse scenarios.

## 7 Acknowledge

We would like to thanks Yu-hui Chen, Yuqi Li, Qifei Wang, Zonglin Li, Zhong Meng, Alex Chen, Jingbin Wang, Abhanshu Sharma, Alec Go, Jesper Andersen, David Racz, DeLesley Hutchins, Luke Sernau, David Munday, and Ewa Dominowska for their suggestions.

## Appendix A Prompts and Instructions

### A.1 Prompts Used for LLMs

We share the prompt used for calling LLMs to get per context assessment in Table 9.

**Table 9: The zero-shot prompts for LLM labeling and critique.**
| Round 1 prompt |
| --- |
| You are now doing a reading comprehension task. It is important that you be as thorough, detail-oriented, and accurate as possible in your response.<br>You are given a question, a set of accepted answers, a document and its title. The document does not necessarily contain the right answer to the question.<br>You should read the title and the document and then check if they provide one of the correct answers to the question.<br>If the title and document together contain the correct answer to the question, output a score of 1.0, otherwise output a score of 0.0.<br>question: <question><br>accepted answers: <answers><br>title: <title><br>document: <document><br>output: |
| Round 2 prompt |
| Your job is to correct another model’s performance on a reading comprehension task.<br>The model was given a question, a set of accepted answers, a document and its title. The document and title do not necessarily contain the right answer. The model was instructed to output a score of 1.0 if the document contains the answer, and a score of 0.0 otherwise.<br>You will be given the same information as the other model along with its output. You should read the title and document and then check if they provide one of the correct answers to the question.<br>Then check if you agree with the previous model’s output.<br>If you agree, output the same score unchanged.<br>If you disagree, output the corrected score.<br>Your output should be as accurate as possible.<br>question: <question><br>accepted answers: <answers><br>title: <title><br>document: <document><br>previous model’s score: <score><br>output: |

### A.2 Rater Guidelines

We share the instructions provided to the human labelers in Table 10.

**Table 10: Instructions for raters creating ground-truth relevance dataset.**
| Human Rater Instructions |
| --- |
| Please read the question, the answer and the context. Please answer if the context can help answer the question. If it can, select 1. Otherwise select 0.<br>1: good<br>0: bad<br>Please use the answers as a hint. However, do not use "is the answer in the context?" as a heuristic for making the decision. |

## Appendix B Dataset Analysis

During the human-labeling process, several raters flagged documents and questions that were difficult to label. In total, 23 out of 500 documents were flagged and 15 out of 50 questions were flagged.

We observed trends in questions and contexts shared in Table 11 that raised concerns about whether and how a human would be able to assess the relevance of the context. These concerns extend to expectations of how well LLMs would be able to do that as well.

**Table 11: Overview of trends, datasets, and examples with associated comments.**
| Trend | Dataset | Example question | Comments |
| --- | --- | --- | --- |
| Questions with time-dependent answers | NQ | who is the president of usa right now | Depends on when the question is asked. |
| NQ | who is the current director of the us mint | Depends on when the question is asked. |  |
| NQ | when is the next deadpool movie being released | Depends on when the question is asked. |  |
| NQ | total number of death row inmates in the us | Fluctuates over time. |  |
| PopQA | What is Prague the capital of? | The borders in this region and the name of the country have changed several times in the 20th century. |  |
| PopQA | What is Dennis Rodman’s occupation? | The accepted answers are "actor, actress, actors, actresses". He was an actor later in his career, but he rose to prominence as a professional basketball player. |  |
| Missing synonyms, abbreviations, and aliases, combined with unclear granularity | NQ | in which regions are most of Africa petroleum and natural gas found | "Region" can refer to different levels of granularity (e.g. Sub-Saharan Africa vs. Ethiopia), but the only accepted answer is "Nigeria". |
| NQ | what type of car is a jeep | The accepted answers are only "off-road vehicles", "light utility vehicles", "sport utility vehicles", but "SUV" is clearly a correct answer as well. |  |
| Non-exhaustive list of answers | NQ | cast of law & order special victim unit | The accepted answers include 16 cast members, but the show went on for 25 seasons with many cast changes and guest stars not included in the list. |
| Oddly phrased question | NQ | right to property according to the constitution of India is a | The only correct answer is "constitutional right", but that is included in the question. It’s not clear what type of answer would be appropriate here. |
| Overly specific answers expected | NQ | where does the story the great gatsby take place | The only accepted answer here is "Long Island of 1922", but the place is Long Island and the question does not ask about when the story is set. |
| Question refers to an entity with a common name without disambiguation | PopQA | What genre is Frances? | There is a musician and a film called "Frances", and both of those could arguably have a genre associated with them. |
| PopQA | Who was the producer of Hurt? | The question is referring to a song performed by Christina Aguilera but there are many other songs, movies, and other entities that share the name and also have a producer. |  |
| PopQA | What is the capital of Cherokee County? | There are many different Cherokee Counties in different states in the USA. |  |

We explored several ways of filtering our human-labeled subsample of Natural Questions to determine how they impacted context assessment F-scores overall and for each auto-rater. We provide two additional filtered versions of the human-labeled RAG relevance dataset as alternatives. See Table 12 for statistics on the filtered datasets and Table 13 for the auto-rater F-scores for each. Both statistical filtering approaches (e.g. removing contexts with non-unanimous labels) and targeted filtering approaches (e.g. removing questions or contexts flagged by human raters) lead to some improvement in F-scores for relevance labels.

**Table 12: Summary of Dataset Evaluations**
| Dataset | # Questions | # Docs | Total % Relevant | Average # Docs per Question | Std Dev # Docs per Question | Average % relevant docs per question | Std Dev % relevant docs per question |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Original 3-rater labels | 50 | 500 | 0.30 | 10 | 0 | 0.30 | 0.23 |
| Revisited (released) | 50 | 500 | 0.31 | 10 | 0 | 0.31 | 0.26 |
| Filter non-unanimous (released) | 50 | 351 | 0.23 | 7.02 | 2.59 | 0.27 | 0.29 |
| Filter flagged docs | 48 | 467 | 0.31 | 9.729 | 0.57 | 0.31 | 0.27 |
| Filter flagged questions | 35 | 350 | 0.29 | 10 | 0 | 0.29 | 0.26 |
| Filter flagged docs and question (released) | 34 | 330 | 0.29 | 9.7 | 0.58 | 0.29 | 0.27 |
| Filter 4-rater split | 50 | 444 | 0.28 | 8.88 | 1.49 | 0.29 | 0.27 |

**Table 13: Evaluation of Labeling Methods w/ Filtering**
| Dataset Filters | # Docs | Rater Model | Critic Model | Average F1 |
| --- | --- | --- | --- | --- |
| Specialized rater corrections | 500 | PALM2 XL | n/a | 0.729 |
| PALM2 XL | PALM2 XL | 0.781 |  |  |
| Gemini Ultra | n/a | 0.761 |  |  |
| Gemini Ultra | Gemini Ultra | 0.704 |  |  |
| PALM2 XL | Gemini Ultra | 0.728 |  |  |
| Gemini Ultra | PALM2 XL 340B | 0.821 |  |  |
| Filter non-unanimous docs | 351 | PALM2 XL | n/a/ | 0.741 |
| PALM2 XL | PALM2 XL | 0.811 |  |  |
| Gemini Ultra | n/a | 0.792 |  |  |
| Gemini Ultra | Gemini Ultra | 0.739 |  |  |
| PALM2 XL | Gemini Ultra | 0.763 |  |  |
| Gemini Ultra | PALM2 XL | 0.856 |  |  |
| Filter flagged docs and questions | 330 | PALM2 XL | n/a | 0.750 |
| PALM2 XL | PALM2 XL | 0.797 |  |  |
| Gemini Ultra | n/a | 0.782 |  |  |
| Gemini Ultra | Gemini Ultra | 0.741 |  |  |
| PALM2 XL | Gemini Ultra | 0.753 |  |  |
| Gemini Ultra | PALM2 XL | 0.833 |  |  |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: Please see weaknesses.
soundness: 3
strengths: 1.	The paper is clearly written and the proposed sparse RAG by appending document assessment stage is well validated. 
2.	In the experiments, the proposed method leads to additional improvements in terms of both RAG and sparse RAG. It is more efficient than RAG, and it is remarkable the sparse RAG shows improved performances over RAG.
summary: The paper proposes a sparse RAG which consists of two stages – 1) document assessment, and 2) generation, to improve the efficiency of RAG. The document assessment selects the relevant documents by individually checking their relevance, and generation part only uses the selected documents as a context to generate an answer, but in an efficient manner where documents are cached, and cross-attention b/w documents is not allowed. To my understanding, the efficiency is obtained here, because all previous information could be cached (no recomputation) after document assessment. Experiment results show that a sparse RAG leads to improvements in terms of both effectiveness and efficiency, i.e., it improves RAG, more efficiently than PCW RAG.
weaknesses: 1.	Some parts in the presentation are not clear. In multitasking data format, more details of the format need to be provided. E.g. What is the information of Rating field? In Figure 1, [ASSESS] and [GENERATE] are special tokens? Detailed position information for cached document representations used in document assessment and generation. Document assessment does not use the cached representations in input?  
2.	Context reduction methods have been widely studied for RAG. Document assessment method needs to be validated separately. Other context compression and reduction baselines (not only RAG) need to be compared.
3.	Given the restricted cross-attention manner across passages, there is concern that the general ability is still well-maintained. SparseRAG needs to be further validated under other advanced CoT, long chain of reasoning, and instruction following tasks. 
4.	Document assessment can also be parallelizable, using restricted cross-attention method? Still, it remains unclear whether the current decision of relevance for documents in an one-by-one manner is optimally designed.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: See the weakness part.
soundness: 3
strengths: 1. The writing is clear and easy to follow, addressing a genuine concern.

2. The proposed method demonstrates efficiency and effectiveness compared to the baselines.

3. Most ablation studies and analyses are convincing.
summary: This paper presents SparseRAG, an efficient framework designed to address the scaling challenges of traditional RAG models. The primary issue tackled is the significant increase in memory usage of the key-value (KV) cache when processing large amounts of retrieved lengthy documents. This issue hurts the ability of large language models (LLMs) to efficiently manage long-form context and capture essential content.

To overcome this, the authors propose SparseRAG, which utilizes a required LLM rater and an optional LLM critic to generate Per-Context-Assessment (PCA) training data. The LLM is then trained on this PCA data to enhance its capability in assessing the relevance of retrieved documents. During both training and inference, after retrieving the related contexts, the LLM first evaluates the relevance of each retrieved documents and simultaneously extracts the KV caches, which are then concatenated to generate the final outputs.

The experimental results indicate that SparseRAG significantly improves both prefill and decoding efficiency. Additionally, it outperforms all other baseline models in generating high-quality outputs. Further analysis and ablation studies confirm that SparseRAG is both an efficient and effective framework.
weaknesses: 1. Ablation Study Suggestion: It would be beneficial to include an ablation study to evaluate the impact of concatenating selected KV caches instead of the corresponding raw contexts. This approach aims to save computation time, but it’s important to assess any potential performance loss compared to recalculating KVs with full attention on concatenated raw contexts. Understanding this trade-off is crucial, as certain applications might accept additional computation time if it results in significant improvements.

2. Comparison with KV Cache Compression Methods: Since the paper claims that SparseRAG enhances efficiency during the prefill and decoding stages by adopting context selection, it would be compelling to compare it with KV cache compression methods such as H2O, SnapKV, and KVMerger. Although these are typically training-free methods, and thus challenging to directly compare, such an analysis could provide valuable insights into the relative benefits and limitations of SparseRAG.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Does reusing the KV Cache during batch inference cause a decline in the quality of responses when the document-level attention scores are missing?
2. Why does the use of an external classifier in the CRAG method increase so much latency? (My understanding is that the classifier should be a model with relatively few parameters, and it should be able to perform batch inference.) Is the training data for the classifier consistent with the PCA task?
soundness: 3
strengths: 1. Proposes an innovative sparsity-driven retrieval method (Sparse RAG) that enhances both inference efficiency and output quality.
2. Successfully demonstrates generalizability and adaptability through evaluations on multiple tasks and datasets.
3. Offers thorough comparative analysis against existing methods, providing clear evidence of computational and quality improvements.
4. Utilizes a practical on-device evaluation setting to substantiate real-world applicability, particularly on resource-constrained devices.
summary: The paper proposes "Sparse Retrieval-Augmented Generation" (Sparse RAG), a paradigm to accelerate large language model (LLM) inference by selectively loading only the most relevant retrieved contexts. This approach aims to improve computational efficiency and generation quality by filtering unnecessary documents during decoding, evaluated across diverse datasets with positive outcomes in both quality and speed.
weaknesses: 1. 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. This aspect could be further explored to strengthen the claims made regarding the effectiveness of the Sparse RAG approach.
2. The training process involves multi-tasking and parallel context encoding, which may introduce additional complexity and require careful tuning to ensure effectiveness.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: The points outlined in the “weakness” should be clarified in the rebuttal.

Also, in your appendix, you mentioned "We observed trends in questions and contexts shared in Table 12 that raised concerns about whether and how a human would be able to assess the relevance of the context. These concerns extend to expectations of how well LLMs would be able to do that as well." 

For example, you mentioned that in the dataset NQ, a query like "who is the president of USA right now", and you commented that "the answer will depend on when the question is asked".  

This is true but I assume that is exactly why RAG could be useful since RAG will retrieve additional information based on the context of the query, and in many cases, information in the query, its surrounding context and the supporting datasets should be able to disambiguate the time period and thus return the correct answer.   I think some discussion on this is necessary.  It would be helpful if the authors may elaborate on how the Sparse RAG approach specifically handles such time-sensitive or ambiguous queries. The authors may include examples or experiments demonstrating how their method performs on queries that require temporal context, or a discussion on how such capabilities can be incorporated in the Sparse RAG framework.
soundness: 3
strengths: 1.	Overall, the paper is clearly written.  

2.	The motivation of SPARSE RAG is natural, and the design of the approach seems to be innovative and reasonable.

3.	The claim that “Sparse RAG combines the assessment of each individual context and the generation of the response into a single process, in which special control tokens are used to prompt the LLM to assess the relevance of each retrieved context, and then only the key-value caches of the most relevant contexts are loaded for decoding using another control token.”  This seems to be interesting and novel, but some detailed explanation is needed to help readers understand the process.  For example, it will be good to explain how the special control tokens work in practice, or to provide a step-by-step explanation of how the context assessment and response generation are combined.
summary: This paper proposes SparseRAG, a method to improve efficiency/effectiveness and reduce computation costs through sparsity. It encodes retrieved documents in parallel, which eliminates latency introduced by long-range attention of retrieved documents. Then, LLMs selectively decode the output by only attending to highly relevant caches auto-regressively.  This is via prompting LLMs with special control tokens. Sparse RAG can facilitate the reduction of the number of documents loaded during decoding for accelerating the inference of the RAG system. It filters out undesirable contexts and enhances the model’s focus on relevant context, improving its generation quality. Evaluations on four datasets show good results on improving generation quality and computational efficiency across multiple tasks.
weaknesses: 1.	The design assumes that a good amount of RAG training data is accessible, which allows effective tailoring and adapting of existing LLMs to the specific needs. But this design relies on pretty heavy training data, for example, in the PoPQA testing, the dataset is split into training, validation and test sets with 8:1:1 ratio (and other testing follows a similar spirit).  I am wondering in practice, an LLM will take many kinds of queries and refer to different kinds of external data, where the training pairs may not be available.  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?  I would suggest that the authors conduct additional experiments with varying amounts of training data to demonstrate the method's robustness. Additionally, it will be useful to discuss potential strategies for adapting the method to domains where extensive training data is not available.

2.	The design does not do any concrete semantic analysis on query intent and nor does any concrete semantic analysis on the text to be retrieved.  It would be helpful to discuss the potential benefits and drawbacks of incorporating semantic analysis into this approach. The reviewer is wondering whether the authors should discuss the potential effect of semantic analysis of the queries and texts and why they chose not to include them in their current design.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. What is the detailed difference between Parallel Context Windows (PCW) and the Parallel Contexts the authors proposed? A specific example would be helpful.
2. In line 203, what is the {Relevant} token?
3. In Figure 1, do $c_{i}$ and $d_{i}$ represent the same thing? If not, what’s the difference between them?
4. Is Parallel Contexts used for both assessment and generation? The authors state that "we unify the generation of the special control tokens and regular vocabulary tokens with one single model," but line 188 mentions that Parallel Contexts is used only for the "generation task." Does this represent a contradiction?
5. In the section *Silver Labels vs. LLM Labels*, could an experiment be added to evaluate Corrective RAG using LLM Labels? The improvement shown in Table 5 appears minimal, which may be influenced by controllable or uncontrollable factors.
soundness: 3
strengths: 1. The proposed framework effectively balances generation quality with computational efficiency.
summary: This work proposes Sparse RAG, a novel RAG framework that reduces computational costs by leveraging sparsity. Sparse RAG encodes retrieved documents in parallel and allows LLMs to selectively decode by attending only to the most relevant cached information. Experiments on four datasets demonstrate that Sparse RAG effectively balances generation quality with computational efficiency.
weaknesses: 1. This work lacks novelty; the authors' primary contribution lies in scoring documents and using an existing framework with modified position encodings to achieve parallel encoding.
2. The experiments lack ablation studies for the module, especially an experiment where only Parallel Contexts is added to RAG (for only generation task).
3. Some parts of this work are not clearly written, making them difficult to understand. For example, lines 201 to 204, with line 204 being especially unclear.

Meta-review: Summary of the paper: This paper introduces Sparse RAG designed to enhance the efficiency and effectiveness of traditional RAG. Sparse RAG addresses the challenges associated with increased memory usage of KV caches when processing lengthy documents, which can impede LLMs in managing long-form contexts effectively. Sparse RAG operates in two main stages: document assessment and generation. In the document assessment phase, relevant documents are evaluated for relevance using a required LLM rater and an optional LLM critic, generating PCA training data. During training and inference, the LLM assesses the relevance of retrieved documents, extracting and caching KV information for efficient processing. By encoding retrieved documents in parallel, Sparse RAG eliminates the latency associated with long-range attention mechanisms, allowing LLMs to decode outputs by selectively attending to only the most relevant cached information. This approach, facilitated by special control tokens, improves computational efficiency and generation quality by filtering out unnecessary contexts. Experimental results show that Sparse RAG significantly enhances both prefill and decoding efficiency while outperforming baselines. 

Strengths of the paper:
- Clear Presentation: The paper is generally well-written and easy to follow, but some sections require further clarification to enhance readability.
- Well-motivated Method: The motivation of Sparse RAG is natural, and the design of the approach is innovative and reasonable.
- Good Empirical Results: Sparse RAG is efficient and effective compared to baselines across multiple tasks and datasets. The authors provide a thorough comparative analysis against existing methods, offering clear evidence of computational and quality improvements, and emphasizing real-world applicability through practical on-device evaluations.

Weaknesses of the paper:
- Missing Ablation Studies (Reviewers p444, ARp8, 5W8k): Include an ablation study to evaluate the impact of concatenating selected KV caches instead of the corresponding raw contexts.
- Needs Further Dicussions on Specific Queries (Reviewer 4HHs): Elaborate on how the Sparse RAG approach specifically handles such time-sensitive or ambiguous queries.
- Missing Comparisons with Corrective RAG (Reviewer pGev): Include experiment to evaluate Corrective RAG using LLM Labels.

Reasons for the decision: 
After considering the rebuttal, I believe the authors have adequately addressed most of the concerns raised. The reviewers discussed these points in detail, and all agree that the paper should be accepted at ICLR. Reviewer pGev rated the paper a 5, which is marginally below the acceptance threshold; however, after further discussion with the authors, they expressed their belief that this work is worthy of acceptance. Overall, Sparse RAG effectively balances generation quality with computational efficiency, marking a promising advancement in the field of retrieval-augmented generation systems.

Decision: Accept (Poster)