论文解读-《REFRAG Rethinking RAG based Decoding》

1. 论文介绍

论文题目：REFRAG: Rethinking RAG based Decoding
论文领域：RAG
论文发表：2025.09 首发于Arxiv
论文代码：https://github.com/facebookresearch/refrag
论文背景：

2. 论文摘要

大型语言模型 （LLM） 在利用广泛的外部知识来增强多轮次和代理应用程序（例如检索增强生成 （RAG））中的响应方面表现出了卓越的能力。然而，处理长上下文输入会引入显着的系统延迟，并需要大量内存来存储键值缓存，从而导致吞吐量降低以及知识丰富和系统效率之间的根本权衡。虽然最大限度地减少长上下文输入的延迟是 LLM 的主要目标，但我们认为 RAG 系统需要专门考虑。在 RAG 中，大部分 LLM 上下文由检索中的串联段落组成，只有一小部分与查询直接相关。由于重新排名过程中的多样性或重复数据删除，这些段落通常表现出较低的语义相似性，导致块对角线注意力模式与标准 LLM 生成任务中的模式不同。基于这一观察结果，我们认为在解码期间对 RAG 上下文进行的大多数计算都是不必要的，并且可以在对性能影响最小的情况下消除。为此，我们提出了 REFRAG，这是一个高效的解码框架，可以压缩、感知和扩展以改善 RAG 应用程序中的延迟。通过利用这种注意力稀疏结构，我们展示了第一个 30.85× 标记的时间加速（对以前的工作的 3.75× 改进），而不会损失困惑。此外，我们针对大上下文的优化框架使 REFRAG 能够将 LLM 的上下文大小扩展为 16× 。我们在各种长上下文任务中对 REFRAG 进行严格验证，包括 RAG、多轮对话和长文档摘要，涵盖广泛的数据集。 实验结果证实，与 LLaMA 模型和其他跨各种上下文大小的最先进的基线相比，REFRAG 提供了显着的加速，而不会损失准确性。此外，我们的实验表明，REFRAG 扩展的上下文窗口进一步提高了流行应用程序的准确性。

3. 相关介绍

3.1 背景介绍

LLM 推理吞吐量会随着上下文的增加而下降，从而限制了它们在需要高吞吐量和低延迟的场景。
增加上下文学习的提示长度会导致推理 过程中更高的延迟和更大的内存消耗 。 具体来说，较长的提示需要额外的键值 （KV） 缓存内存，该缓存随提示长度线性扩展。此外，首个token时间 （TTFT） 延迟呈二次方增加，而迭代token时间 （TTIT） 延迟随提示长度 线性增长。

优化具有大范围上下文的LLM的推理延迟，是一个活跃的研究领域，方法从修改注意力机制的复杂性到稀疏注意力和上下文， 以及改变环境喂养策略。
本文重点关注基于RAG的应用，目标是改善推理延迟，特别是TTFT。本文认为利用RAG上下文固有的独特结构和稀疏性的专业技术可以大大减少内存和计算开销。将RAG TTFT视为通用的LLM推理问题忽略了几个关键方面，（1）token分配效率低，（2）编码信息的浪费，（3）异常结构和稀疏的注意力。

3.2 本文贡献

提出了REFRAG （REpresentation For RAG），一种在 RAG 中有效解码上下文的新机制。REFRAG 显着降低了解码过程中的延迟、TTFT 和内存使用，所有这些都不需要修改 LLM 架构或引入新的解码器参数。
在解码过程中的优化：
（1），REFRAG 不使用检索到的段落中的标记作为输入，而是利用预先计算的压缩块嵌入作为近似表示，将这些嵌入直接输入解码器。好处有

• 缩短了解码器的输入长度，提高了代币分配效率
• 可以在检索中重用预先计算的块嵌入，从而消除冗余计算
• 降低了注意力计算的复杂性，注意力计算的复杂性现在随着块的数量而不是上下文中的标记数量而呈二次方扩展。
  （2），REFRAG 支持在任意位置压缩令牌块，同时保留解码器的自回归性质，从而支持多轮和代理应用程序。
  （3），轻量级强化学习 （RL） 策略进一步增强了这种“随处压缩”功能，该策略有选择地确定何时需要完整的块标记输入，以及何时低成本、近似的块嵌入就足够了

4. 模型架构

令Mdec为编码模型，Menc为解码模型。给定一个输入有T个tokens（x1，x2，xT），假设前面q个tokens是主要输入tokens，后面s个token是上下文tokens（在RAG中检索得到）。

整个REFRAG的架构设计为

该模型由仅解码器的基础模型和轻量级编码器模型组成，上下文被分为多个k大小的${C_1, C_2, C_L}$块，$L=\frac{s}{k}$其中$C_i={x_{q+ki}, ..., x_{q+ki+k-1}}$，然后编码器模型处理所有块以获得每个块的块嵌入 $c_i=M_{enc}(c_i)$，该块嵌入与投影层一起投影，以匹配解码器模型的token嵌入大小，这些投影的块嵌入与问题的标记嵌入一起提供给解码器模型。

延迟和吞吐量改进
基于三个指标，TTFT首个token的延迟，TTIT迭代生成token的延迟，吞吐量是单位时间生成的tokens数量。
对于较短的上下文长度，我们的方法实现 k× 的TTFT 和吞吐量的加速。对于较长的上下文长度，两个指标的加速速度最高可达 。

5. 方法

为了对齐编码器和解码器，使用了CPT（持续预训练）算法，对于每个数据数据点，它包含 s+o=T 数量的标记，我们使用 CPT 来准备模型，以利用块嵌入为下游任务准备模型。为了进一步提高性能，我们通过 RL 引入了选择性压缩。通过 CPT 对齐编码器和解码器后，我们应用监督微调（SFT）使模型适应特定的下游任务，例如 RAG 和多轮对话。

在 CPT 期间，我们将第一个 s 标记输入编码器 $x_{1:s}$ ，并使用其输出来协助解码器预测下一个 o 标记 $x_{s+1:s+o}$。这项任务鼓励模型利用上下文信息进行下一段预测，从而为下游应用做好准备。目标是对齐任何编码器-解码器组合，以便使用压缩上下文生成的代与原始解码器生成的代非常相似，并可以访问完整上下文。

5.1 持续预训练方法

重建任务
在这个任务中，我们冻结解码器模型，只训练编码器和投影层 。主要目标是对齐编码器和投影层，以便：1）编码器可以压缩 k 令牌，同时将信息丢失降至最低，2）投影层可以有效地将编码器的块嵌入映射到解码器的标记空间中，使解码器能够解释并准确地重建原始信息。专门选择重建任务是为了鼓励模型在训练过程中依赖上下文记忆而不是其参数记忆。一旦编码器通过此重建任务与解码器对齐，我们就会通过解冻解码器来启动 CPT。

5.2 课程学习

随着块长度 k 的增加，可能的标记组合数量呈指数级增长，特别是以 $V^k$ 的速率扩展，其中 V 是词汇量。在固定长度的嵌入中有效地捕获这种多样性是一项重大挑战。此外，从块嵌入重建 s=k×L，L的token进一步增加了任务的难度。

与直觉相反的是，直接继续对解码器进行预训练以利用编码器输出并没有减少困惑，即使对于重建任务也是如此。

5.3 选择性压缩

REFRAG 引入了选择性标记压缩，扩展了未压缩的重要上下文块，以改进答案预测。RL 策略以下一段预测困惑度作为负奖励为指导，确定哪些块要保留其原始形式。编码器和解码器经过微调，可以处理压缩和未压缩块的混合输入。策略网络利用块嵌入和屏蔽来优化顺序块扩展，从而保留解码器的自回归属性并实现压缩的灵活放置。

6. 实验设置

训练数据集选择
选择Slimpajama 数据集，是一个用于 LLM 预训练的开源数据集。该数据集包含来自 Wikipedia、Arxiv、Books、StackExchange、GitHub、Commoncrawl、C4 的数据。我们只使用数据集中的 Book 和 ArXiv 域，因为这两个域包含长文本

评估数据集：检查模型的泛化，我们还报告了 PG19 和 Proof-pile 数据集的结果

基线模型：所有基线模型均基于 LLaMA-2-7B

实验结果

消融实验
课程学习对于重建任务的有效训练至关重要

重建任务对于模型学习持续的预训练任务至关重要。

基于 RL 的选择性压缩的优点
在各种压缩速率下，通过改变要压缩的块数（即调整）来实现 p

REFRAG 在不同压缩率下训练
说明了 REFRAG 在连续预训练任务中不同压缩率下的训练轨迹。我们观察到随着压缩率的增加而性能下降

7. 上下文学习的应用

本节将研究对从预训练阶段获得的模型进行微调，以解决各种下游任务，包括 RAG、长文档摘要和与 RAG 的多轮对话。对于每个应用程序，我们策划了一个指令调整数据集，以促进模型微调。

7.1 检索增强生成

采用的训练数据集
对话 ：OpenAssistant 对话数据集。 开放领域 QA：CommonsenseQA、MathQA、 网络问题、Wiki 问答、Yahoo！ Answers QA、FreebaseQA、MS MARCO。 阅读理解 ：段落离散推理，PubMedQA，QuaRel，SQuADv2。 思维链推理 ：代数 QA 及其基本原理、CommonsenseQ 解释、小学数学 8K、MathQA、StrategyQA。

评估数据集
训练数据集中每个数据集的 5% 的数据保留下来进行评估。
在 2 个设置上评估我们的性能：1） 强检索器 ：在此设置中，我们使用强检索器并检索 K 最近邻来回答问题;2） 弱检索器 ：在此设置中，我们检索 200 个段落并随机选择 K 个段落来回答问题。弱检索器设置与现实世界的系统非常相似，因为 RAG 检索系统经常遭受子系统之间的错误累积。

7.2 多轮对话

三个不同的知识密集型多轮对话数据集：TopiOCQA 、ORConvQA和 QReCC。

实验结果

结果显示了不同数量的对话回合和检索到的段落的结果。我们的模型在 5 次传代设置中的三个数据集中的两个 LLaMAFT 以及 10 次传代设置中的所有三个数据集上都表现出色。这种改进归因于 LLaMAFT 有限的 4k token 上下文窗口，这需要在较长的上下文中截断部分对话历史记录，从而导致丢失回答后续问题所需的关键信息。相比之下，由于压缩方法的优势，我们的模型在相同的 LLaMA 模型上训练，没有扩展其有效的位置编码，即使在大量段落的情况下也能保持稳健的性能。

8. 相关工作

8.1 检索增强语言建模

改进检索增强生成的方式有多种，（1）引入了检索增强掩码语言模型的预训练；（2）工作侧重于通过在检索到的段落前面插入提示并采用集成方法生成响应来微调现有的LLM；（3）引入了融合解码器，它使用编码器并行处理每个段落，并通过解码器连接隐藏状态以生成。

8.2 高效的长上下文LLM（Efficient Long-Context LLMs）

各种策略来减少 LLM 长上下文生成中的内存使用和加速延迟。有（1）压缩注意力，将注意力复杂度从二次降低为线性；（2）提出了注意力接收器来减少长上下文生成的 KV 缓存，尽管这并不能减少预填充阶段的延迟；（3）采用交叉关注来自上下文token的标记嵌入，减少了 KV 缓存和注意力计算。
CEPE 仅限于前缀上下文应用，因为它破坏了上下文的因果结构，使其不适合多轮 RAG 或摘要等任务。此外，CEPE 不使用令牌压缩，导致解码延迟相似甚至增加。

8.3 压缩Transformer

方法有多种
（1）引入了压缩转换器，它压缩 KV 缓存以减少长上下文应用程序的内存需求。然而，这种方法只会减少 KV 缓存的使用量，不会改善第一个token延迟的时间，并且需要从头开始训练模型。
（2）通过采用递归上下文压缩来扩展这一想法，为每个块生成一个汇总隐藏状态，以通知下一个块的计算。然而，递归性质会阻止块嵌入的预计算和重用，并且不会减少解码延迟。

8.4 提示压缩

提示压缩旨在减少输入令牌长度，以降低延迟和成本，同时保持任务性能。一种突出的方法是 LLMLingua， 它采用从粗到细、预算控制的压缩和令牌级迭代细化，以最小的性能损失实现高压缩比。

9. 总结

REFRAG算法是使用压缩算法，增大上下文的范围，综合了现有的优化手段来实现的较大改进，在减少内存使用和加快推理上都有较大提升。

10. 个人感想

开始入门RAG，对整个RAG的研究了解不多。大概可以看到更大的上下文，更快的推理速度是这个领域的当前研究方向。更大的上下文范围意味着LLM回复更加充分概括和包含信息，答案更加完善。更快的推理速度指的是RAG的更加高效。