Xiaoqiang Lin Aritra Ghosh Bryan Kian Hsiang Low Anshumali Shrivastava Vijai Mohan

摘要

大型语言模型（LLMs）在利用大量外部知识以增强多轮交互和代理型应用（如检索增强生成，RAG）中的响应能力方面展现出卓越性能。然而，处理长上下文输入会引入显著的系统延迟，并对键值缓存（key-value cache）产生巨大的内存需求，导致吞吐量下降，从而在知识丰富性与系统效率之间形成根本性权衡。尽管降低长上下文输入的延迟是LLM的核心目标之一，我们认为RAG任务需要特殊的考量。在RAG中，模型上下文主要由检索结果拼接而成，其中仅有少量片段与查询直接相关。这些片段由于重排序过程中的多样性或去重处理，往往具有较低的语义相似性，从而导致注意力机制呈现块对角结构，这与标准LLM生成任务中的注意力模式显著不同。基于这一观察，我们认为在解码过程中对RAG上下文进行的大部分计算是冗余的，可通过有效剪枝实现，且对性能影响极小。为此，我们提出REFFRAG——一种高效的解码框架，通过“压缩-感知-扩展”机制显著提升RAG应用的推理速度。通过利用上下文中的稀疏结构特性，我们在不牺牲困惑度（perplexity）的前提下，实现了首次词生成时间（time-to-first-token）30.85倍的加速（相较此前工作提升3.75倍）。此外，我们的优化框架使LLM能够将上下文长度扩展16倍。我们在多种长上下文任务上对REFFRAG进行了严格的验证，涵盖RAG、多轮对话以及长文档摘要，覆盖广泛的数据集。实验结果表明，相较于LLaMA系列模型及其他当前最先进的基线方法，REFFRAG在不同上下文长度下均实现了显著的加速，同时保持了完全一致的准确性。