RAG 技术通过结合外部知识库,有效提升了大语言模型(LLM)的性能。然而,RAG 系统在处理长序列生成时面临严重的延迟和效率挑战。最近,北京大学联合字节跳动的研究人员提出了一项名为 RAGCache[1] 的创新研究为这一难题提供了巧妙的解决方案。 关键发现
现有的优化方案主要聚焦于LLM推理加速本身,未能充分考虑RAG系统的特点,研究团队通过分析,揭示了三个重要发现:
- 性能瓶颈:RAG系统的主要瓶颈在于LLM生成步骤,注入的外部知识文档显著延长了处理序列。实验数据显示,注入的外部知识文档平均长度达3718个token,远超原始请求的348个token,而计算开销与序列长度呈正比,因此是导致性能显著下降的原因。
- 访问模式:检索请求呈现高度集中的特征,约3%的热门文档占据了60%的检索请求。这种模式为缓存优化提供了重要依据。
- 优化空间:通过缓存文档的中间计算状态,可以显著降低计算开销。实验表明,这种方法可将预填充延迟降低11.5倍。
RAGCache 的核心创新
RAGCache 的核心思想是通过多层动态缓存系统,高效缓存和重用检索文档的中间状态,从而显著提升 RAG 系统的性能。
其主要创新点包括:
- 知识树结构 : RAGCache 设计了一种知识树结构来组织缓存的文档状态。这种树状结构巧妙地处理了 RAG 系统中文档检索的顺序敏感性问题。
例如,考虑两个文档序列:[D1, D3]和[D2, D3]。尽管 D3 在两个序列中都出现,但由于前面文档的不同,其 key-value 张量在各自序列中的值是不同的。知识树结构能够高效管理这些变化,确保快速检索的同时保持文档顺序。
2. 前缀感知的贪心双重大小频率(PGDSF)替换策略 : RAGCache 引入了一种复杂的缓存替换策策略,综合考虑了多个因素:
这种策略确保最有价值的文档状态能够留在缓存中,最大化命中率并最小化重复计算。
-
文档顺序
-
key-value 张量的大小
-
访问频率
-
访问时间的新近性
- 多层缓存 : RAGCache 实现了一个跨 GPU 和主机内存的分层缓存系统。这种方法允许高效利用更快但容量有限的 GPU 内存,同时利用主机内存的更大容量来存储较少访问的状态。
例如,频繁访问的文档状态可能存储在 GPU 内存中以实现快速访问,而不常用的状态则保存在主机内存中。这种分层方法同时优化了速度和容量。
- 动态推测流水线 : RAGCache 最具创新性的特征之一是其通过动态推测流水线重叠向量检索和 LLM 推理的能力。
具体工作流程如下:
- 随着向量搜索的进行,RAGCache 将中期结果发送给 LLM 进行推测生成。
- 如果检索到的文档发生变化,系统会适应性地启动新的推测生成。
- 这种方法最小化了空闲时间,可以显著减少端到端延迟。
性能提升
RAGCache 取得的结果令人印象深刻:
- 与 vLLM+Faiss 相比,首个 token 生成时间(TTFT)最高减少 4 倍
- 与 vLLM+Faiss 相比,吞吐量最高提升 2.1 倍
- 与 SGLang 相比,TTFT 最高减少 3.5 倍
- 与 SGLang 相比,吞吐量最高提升 1.8 倍
这些性能提升在不同模型和检索设置下都保持一致,展示了 RAGCache 方法的稳健性。
实际应用
RAGCache 的改进对广泛的应用有重要影响:
- 问答系统 : 对于需要外部知识的复杂查询,可以实现更快的响应时间。
例如,在一个客户服务聊天机器人系统中,当用户询问"如何申请退款"时,RAGCache 能够快速从缓存中检索相关政策信息,大大缩短响应时间。 2. 内容生成 : 更高效地创建上下文感知的内容,适用于聊天机器人或自动写作助手等应用。
比如,在新闻摘要生成任务中,RAGCache 可以快速检索和缓存相关背景信息,使得生成的摘要更加准确和全面。 3. 研究工具 : 为处理大型数据集的研究人员和分析师提供更快的相关信息访问。
在科研领域,RAGCache 可以加速文献综述过程,研究人员能够更快地检索和比较大量相关论文的关键信息。
结论
RAGCache 代表了优化检索增强生成系统的重要进展。通过解决长序列生成和计算开销的核心挑战,它为更具响应性和效率的 AI 应用铺平了道路,使得先进的 AI 能力在实际应用中更易于实现和使用。
随着 LLM 的不断发展和知识密集型任务需求的增长,像 RAGCache 这样的系统将在使先进 AI 能力更易获取和实用化方面发挥关键作用。
参考资料
[1] RAGCache: https://kiarxiv.org/pdf/2404.12457
后台回复“进群”入群讨论