TurboRAG: Accelerating Retrieval-Augmented Generation with Precomputed KV Caches for Chunked Text

We would like to express our gratitude to the reviewer for your valuable feedback.

While the idea of precomputing the KV cache to reduce TTFT is effective, it is not particularly novel. Prior work, such as [1], has already explored precomputing KV caches to achieve similar objectives. Although this paper mentions RAGCache, it lacks a thorough discussion that differentiates the proposed methods from existing approaches. Including a detailed comparison with RAGCache in both methodology and experimental results would strengthen the paper’s contribution.

RAGCache is quite different from our approach, as it stores complete key-value caches of identical queries that frequently appear in historical dialogue records (corresponding to contexts formed by the same recalled document chunks). Therefore, compared to TurboRAG, it faces two challenges. First, if the order of the recalled documents differs, RAGCache cannot be used directly. In contrast, TurboRAG has additional processing for position embeddings, allowing it to handle variations in the order of recalled documents. Second, RAGCache has a hit rate issue; it can only recall contexts that exactly match those in the index along with the prompt text, requiring recalculation for any discrepancies. TurboRAG does not face this issue, ensuring a consistent speedup. We will include this discussion in the related work section.

While precomputing document KV caches can effectively reduce TTFT, it increases storage costs, as each document needs a separate set of KV caches for different models. It is important for the paper to mention this storage issue so that readers can understand the potential trade-offs involved.

For the storage issue, please refer to the overall comments C1.

We appreciate that the reviewer understands and recognizes the contributions of this work.  We address the main concerns as follows.

It is unclear whether the efficiency improvements are significant when the retrieved documents are short (efficiency improvements are only measured with 8-16k tokens. I believe that most RAG settings will be mostly working with shorter prompts than that. For example, the supporting passages in multi-hop QA datasets within LongBench are actually only around 200 tokens each.

Please refer to the overall comments C2. Based on your feedback, we have added the experimental data, which can be found in Table 7 and Table 9 in our new version. If each text chunk contains 200 tokens, recalling and concatenating 5 segments results in 1000 tokens. According to the experimental results in Table 9, even with a batch size of 1, we can achieve a commendable speedup of up to two times.

The metrics presented in the results sections are not well specified.

Are you referring to the LongBench scores? We have added an explanation regarding this aspect.

The table column and row titles are, especially for Table 2 and 3, need to be properly capitalized and formatted.

Thank you very much for your reminder; we have addressed this issue.

Did you have any experiments on using the KV cache representations to enhance retrieval itself? It seems awkward to have to encode all documents with two models offline.

This is a very interesting idea, as it could reduce the number of models. A potential issue is that the cost of calculating similarity using the LLM's KV cache can be quite high. We forsee a signficant improvement in retrieval if the similarity computation overhead can be reduced.

Is there any reason why you would expect the Naive RAG model would underperform Turbo reordered RAG model in english LongBench QA tasks? This seems very strange to me.

Thank you for the comment. One possible reason that the test set recalls many documents, but many of them are unrelated. The computation of cross-attention may lead to mutual interference, which therefore leads to the accuracy drop.

The fact that HotpotQA and DuReader are part of the fine-tuning data to enable TurboRAG and also the main experimental setting is makes it harder to tell if the method can truly generalize. Are we at least fully certain that there is no data leakage?

Please refer to the overall comments C4.

Thank you for the thoughtful review of our work! Please allow us to address your concerns and answer the questions.

Insufficient related work discussion

Thanks for the suggestion. We have added related work on sparse attention and reordered positional IDs in the revised version.

Limitations of model fine-tuning

Thanks for pointing out. We agree that fine-tuning limits the applicability of our work and prevents it from being directly used by newly emerging state-of-the-art LLMs. We are currently exploring ways to reduce or even eliminate the dependency on fine-tuning. However, the majority of existing techniques that attempt to extend the context window length have to modify positional embedding. In order to adapt to the increased context window length, these techniques commonly tend to continue pretraining or fine-tune the foundatoin model. Inpired from these works, we hence employed fine-tuning to guarantee that TurboRAG delievers "the best" model performance.

Storage cost

Please refer to overall comment C1. We have also included this part in the discussion of limitations. Thank you for your reminder.

Documents integration performance

We conducted additional experiments on the multi-document fusion test set, RGB Integration, which contains answers across multiple documents and effectively tests the model's information fusion capability. The experimental results indicate that without fine-tuning, whether using the composite or reordered approach, directly filling the document KV Cache yields poor performance. After fine-tuning, even with strong correlations between documents, the performance of TurboRAG is comparable to that of Naïve RAG.

Limitation section

Thank you for your reminder; we have added a discussion on limitations in the new version.

GPT-4o Chinese performance and TTFT in RGB

Our test set includes both English and Chinese data, and GPT-4o has strong Chinese capabilities, outperforming many open-source models and ranking first on some authoritative Chinese benchmarks, such as https://github.com/jeinlee1991/chinese-llm-benchmark. Additionally, our intention was not to prove that our model is superior to GPT-4o, but rather to provide a benchmark reference for the model's RAG capabilities.

Regarding the RGB English dataset, we have added experimental data, which still shows 2.4x speedup over the Naive RAG. We will include the detailed experiments in the appendix.

Why is gpt-4o's baseline not reported for the LongBench multi-doc QA datasets?

The experiments with TurboRAG primarily demonstrate that using the TurboRAG approach on the same model offers latency advantage over the standard RAG method, without compromising the question-and-answer performance. Given limited resources, we only used a 7B model for the experiments, and our intention was not to compare the model's capabilities with GPT-4o.

HotpotQA data leakage

Thanks for the comment. Please refer to overall comment C4.

Experiments for Llama3.1 8B and 70B.

We appreciate the suggestion. Please refer to the overall Comment C3.

Link to code

We have created an anonymous GitHub repository and uploaded our code. The results presented in the paper can be reproduced by following the guideline provided in the repo. (https://anonymous.4open.science/r/TurboRAG-D732)

limitation section

Based on your feedback, we have added a limitations section.

We sincerely appreciate the reviewer's constructive comments and positive acknowledgment of our work. In response to your suggestions, we have expanded the discussion in the related work section.

Limited exploration of trade-offs: The paper focuses primarily on the benefits of TurboRAG but does not thoroughly explore potential drawbacks or limitations. For instance: The authors do not discuss the storage requirements for precomputed KV caches, which could be substantial for large document collections. There's no analysis of how TurboRAG might impact memory usage during inference, especially for scenarios with many retrieved documents. A more balanced discussion of these trade-offs would provide a clearer picture of TurboRAG's applicability in different settings.

For the storage issues, please refer to the overall comments C1. Additionally, we have included a discussion of the limitations in the revised version.

Limited comparison to other optimization techniques: The paper primarily compares TurboRAG to a naive RAG implementation. However, it doesn't extensively compare against other recent optimization techniques for RAG systems or long-sequence processing, such as Efficient attention mechanisms (e.g., Performer, Reformer) and Other caching strategies or optimization approaches for RAG systems. A broader comparison to other RAG optimization approaches in addition to native RAG and also to other LLM architectures would help contextualize TurboRAG's contributions within the current state of the art.

Thank you for your suggestion; we will include this discussion in the related work section.

Limited discussion of scalability: The paper demonstrates impressive speedups, but doesn't extensively discuss how TurboRAG scales with increasing document collections at scale or query complexity. Additional experiments or analysis on scalability would strengthen the paper's claims about TurboRAG's broader applicability.

The longbench test set we are using typically recalls 20-30 documents, averaging around 10,000 tokens, which is quite substantial. For query complexity, your concern is valid, and we will collect datasets with complex queries in our future work to assess their impact on the model. Query complexity is indeed likely to impact the accuracy of the recall, as it tests the capabilities of the recall model.

Could the authors provide an analysis of the storage requirements for TurboRAG compared to conventional RAG systems? This information would help readers understand the trade-offs involved.

For the storage issues, please refer to the overall comments C1.

How does TurboRAG impact memory usage during inference, especially for scenarios with many retrieved documents? An analysis of memory consumption would provide a more complete picture of the system's efficiency.

During inference, loading the KV cache requires additional memory; however, we only load the KV cache for a few document chunks at a time, which keeps memory consumption quite low. As we mentioned earlier, the KV cache for a single document chunk occupies approximately 28MB of storage. Nonetheless, this is an good suggestion, and we will include a discussion on memory consumption in the relevant section.

How does TurboRAG compare to other recent optimization techniques for RAG systems or long-sequence processing, such as efficient attention mechanisms (e.g., Performer, Reformer) or other RAG caching and optimization strategies?

The Performer and Reformer introduce new attention algorithms aimed at reducing the time or space complexity of attention calculations, serving as general optimizations for the Transformer architecture. In contrast, our goal is to reduce the online computation time of LLMs in the RAG scenario (where multiple document chunks are concatenated) by leveraging a space-for-time approach, without making structural optimizations to the LLM itself. However, we agree that it is indeed important to include this comparison in the related work section. Thank you for your valuable suggestion.

The paper focuses on a specific LLM architecture. Have the authors tested or considered how TurboRAG might apply to other popular LLM architectures? This information would be valuable for understanding the broader applicability of the approach.

Please refer to the overall comments C3.

Thank you for your continued response. We have uploaded a new version of the paper to this platform, which includes updates on related work, a discussion of limitations, and some additional experimental results.

Regarding other model architectures, we are currently training TurboRAG on LLaMA 3.1, and we will add the experimental data to this platform as soon as possible.