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

📄 arXiv: 2410.07590v1 📥 PDF

作者: Songshuo Lu, Hua Wang, Yutian Rong, Zhi Chen, Yaohua Tang

分类: cs.CV, cs.CL

发布日期: 2024-10-10

💡 一句话要点

TurboRAG：通过预计算KV缓存加速分块文本的检索增强生成

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 检索增强生成 RAG KV缓存 预计算 低延迟 大型语言模型 推理加速

📋 核心要点

  1. 现有RAG系统在预填充阶段处理大量文档块，导致计算开销大，TTFT延迟高。
  2. TurboRAG通过预计算并存储文档的KV缓存，在推理时直接检索，避免在线计算，从而加速生成。
  3. 实验表明，TurboRAG在保持性能的同时，显著降低了TTFT，最高可达9.4倍的加速效果。

📝 摘要（中文）

现有的检索增强生成（RAG）系统为了预填充，需要连接和处理大量检索到的文档块，这需要大量的计算，从而导致显著的time-to-first-token (TTFT)延迟。为了减少计算开销和TTFT，我们提出了TurboRAG，一种新颖的RAG系统，它重新设计了当前RAG系统的推理范式：首先离线预计算并存储文档的键值（KV）缓存，然后直接检索保存的KV缓存以进行预填充。因此，在推理期间消除了KV缓存的在线计算。此外，我们还提供了对掩码矩阵和位置嵌入机制的一些见解，并微调了预训练语言模型，以保持TurboRAG的模型准确性。我们的方法适用于大多数现有的大型语言模型及其应用，而无需修改模型和推理系统。在一系列RAG基准测试中的实验结果表明，与传统的RAG系统相比，TurboRAG将TTFT降低了高达9.4倍（平均8.6倍），同时保留了与标准RAG系统相当的性能。

🔬 方法详解

问题定义：现有RAG系统在推理阶段需要对检索到的文档块进行在线计算KV缓存，这导致了较高的计算开销和较长的time-to-first-token (TTFT)。尤其是在处理大量文档块时，这个问题更加突出。因此，如何降低RAG系统的TTFT，提高推理效率是一个重要的挑战。

核心思路：TurboRAG的核心思路是利用预计算的思想，将文档的KV缓存离线计算并存储起来。在推理时，直接检索预先计算好的KV缓存，避免了在线计算的开销，从而显著降低TTFT。这种方法类似于缓存机制，通过空间换时间，提高了推理效率。

技术框架：TurboRAG的整体框架主要包含两个阶段：离线预计算阶段和在线推理阶段。在离线预计算阶段，系统对文档进行分块，并计算每个块的KV缓存，然后将这些KV缓存存储起来。在在线推理阶段，系统首先检索相关的文档块，然后直接从存储中加载对应的KV缓存，用于预填充，最后生成最终的输出。整个过程无需在线计算KV缓存。

关键创新：TurboRAG最关键的创新点在于将预计算的思想引入到RAG系统中，通过离线计算KV缓存，避免了在线计算的开销，从而显著降低了TTFT。与传统的RAG系统相比，TurboRAG无需在推理时进行KV缓存的计算，而是直接使用预先计算好的缓存，这大大提高了推理效率。

关键设计：TurboRAG的关键设计包括：文档分块策略（如何将文档分割成合适的块），KV缓存的存储格式（如何高效地存储和检索KV缓存），以及如何处理掩码矩阵和位置嵌入，以保证模型的准确性。此外，论文还对预训练语言模型进行了微调，以进一步提高TurboRAG的性能。

🖼️ 关键图片

img_0

📊 实验亮点

TurboRAG在一系列RAG基准测试中表现出色，与传统的RAG系统相比，TTFT降低了高达9.4倍（平均8.6倍），同时保持了与标准RAG系统相当的性能。这表明TurboRAG在加速RAG推理方面具有显著优势。

🎯 应用场景

TurboRAG适用于需要快速响应的RAG应用场景，例如在线问答、对话系统、信息检索等。通过降低TTFT，TurboRAG可以显著提升用户体验，并降低计算成本。未来，该技术可以应用于更广泛的LLM应用，例如智能客服、内容生成等。

📄 摘要（原文）

Current Retrieval-Augmented Generation (RAG) systems concatenate and process numerous retrieved document chunks for prefill which requires a large volume of computation, therefore leading to significant latency in time-to-first-token (TTFT). To reduce the computation overhead as well as TTFT, we introduce TurboRAG, a novel RAG system that redesigns the inference paradigm of the current RAG system by first pre-computing and storing the key-value (KV) caches of documents offline, and then directly retrieving the saved KV cache for prefill. Hence, online computation of KV caches is eliminated during inference. In addition, we provide a number of insights into the mask matrix and positional embedding mechanisms, plus fine-tune a pretrained language model to maintain model accuracy of TurboRAG. Our approach is applicable to most existing large language models and their applications without any requirement in modification of models and inference systems. Experimental results across a suite of RAG benchmarks demonstrate that TurboRAG reduces TTFT by up to 9.4x compared to the conventional RAG systems (on an average of 8.6x), but reserving comparable performance to the standard RAG systems.