REFRAG: Rethinking RAG based Decoding

大语言模型（LLMs）在利用广泛的外部知识来增强多轮对话和智能代理应用（例如检索增强生成，RAG）中的响应质量方面表现出色。然而，处理长上下文输入会带来显著的系统延迟，并需要大量内存来存储键值缓存（key-value cache），从而降低了吞吐量，并在知识丰富性与系统效率之间形成了一种根本性的权衡。虽然减少长上下文输入带来的延迟是大语言模型的主要目标之一，但我们认为RAG应用需要特别的考量。 在RAG中，大语言模型的上下文通常由检索出的多个段落拼接而成，其中只有少数部分与当前查询直接相关。这些段落由于在重排序过程中进行了多样性增强或去重处理，往往语义相似度较低，导致注意力机制呈现出块对角线模式，这与标准的大语言模型生成任务中的注意力模式不同。基于这一观察，我们认为在解码过程中对RAG上下文的大多数计算是不必要的，可以在几乎不影响性能的前提下予以消除。 为此，我们提出了REFRAG，一个高效的解码框架，通过压缩、感知和扩展三个阶段来改善RAG应用中的延迟问题。借助对上下文稀疏结构的利用，我们实现了首词生成时间（time-to-first-token）加速30.85倍（相比之前的工作提升了3.75倍），且没有造成困惑度（perplexity）的损失。此外，我们针对长上下文的大规模优化框架使REFRAG能够将大语言模型的上下文长度扩展至原来的16倍。 我们在多种长上下文任务中对REFRAG进行了全面验证，包括RAG、多轮对话和长文档摘要等场景，覆盖了多个数据集。实验结果表明，在不同上下文长度下，与LLaMA系列模型及其他最先进的基线方法相比，REFRAG在不损失准确率的前提下实现了显著的加速效果。

论文试图解决大型语言模型（LLMs）在检索增强生成（RAG）应用中处理长上下文输入时所带来的系统延迟和内存消耗问题。由于RAG中的上下文通常由多个检索到的段落拼接而成，其中只有少数段落与当前查询相关，因此传统的注意力机制在这些上下文上进行计算时存在大量冗余，影响推理效率。这是一个在RAG应用场景中亟需解决的新问题。

论文提出REFRAG，一个针对RAG场景的高效解码框架。其核心思想是利用RAG上下文中注意力模式的稀疏性（即块对角结构），通过压缩、感知和扩展机制（compress-sense-expand）来减少不必要的注意力计算，从而提升解码速度。相比传统方法，REFRAG专注于识别并保留对当前生成最相关的上下文信息，避免在无关段落上浪费计算资源。

1. 在保持语言模型困惑度（perplexity）不变的前提下，实现了30.85倍的首字生成时间加速（Time-to-First-Token），相比此前工作提升了3.75倍。 2. 通过优化长上下文处理机制，成功将LLM的上下文长度扩展至原来的16倍。 3. 在多个RAG、多轮对话和长文档摘要任务中进行了广泛验证，覆盖多个数据集，验证了方法的通用性和有效性。 4. 实验设计严谨，涵盖了不同上下文长度下的性能对比，但论文未提及是否开源代码。

1. Retrieval-Augmented Generation (RAG) for Language Models 2. Efficient Attention Mechanisms for Long Contexts in Transformers 3. Sparse Attention and Memory Compression in Large Language Models 4. Context Selection and Pruning Techniques in Retrieval-Augmented Models 5. LLaMA系列模型及其在长上下文任务中的优化研究