大型语言模型 (LLM) 改变了我们与信息交互的方式,但它们有一个众所周知的致命弱点: 对计算资源的巨大需求。这一点在 检索增强生成 (Retrieval-Augmented Generation, RAG) 系统中尤为明显,因为这类系统需要将大量外部文本注入模型以帮助其回答问题。我们提供的上下文越多,潜在的答案质量就越高,但处理过程也会变得更慢、更昂贵。这在 知识量效率 之间造成了一个令人沮丧的权衡。

想象一下,在一个由 RAG 驱动的聊天机器人中提问。在后台,系统可能会检索十个不同的文档,将它们拼接成一个巨大的提示,然后发送给 LLM。模型必须阅读并处理该上下文中的每一个词,然后才能生成答案的第一个 token。这种初始处理延迟——即 首个 token 生成时间 (Time-to-First-Token, TTFT) ——是让 RAG 系统感觉真正具有交互性的主要瓶颈。

Meta 的研究人员在近期发表的论文 “REFRAG: Rethinking RAG based Decoding” 中直击这一问题。他们认为,我们一直像对待其他 LLM 任务一样处理 RAG 推理,这是一种浪费。RAG 的上下文不仅不是一个长篇叙述,而是由多个通常不相关的文档组成,这些文档之间的关系高度稀疏。通过利用这种独特的结构,他们开发了 REFRAG——一个解码框架,可将 RAG 系统的 TTFT 提速高达 30.85 倍,且不牺牲准确性。

在这次深入探讨中,我们将解析 REFRAG 背后的核心思想——它如何压缩、感知和扩展上下文——并探讨让它如此高效的训练策略和实验结果。


RAG 中“低效”上下文处理的问题

这种低效率始于标准 LLM 处理长提示的方式。在自注意力机制中,每个 token 都会与其他所有 token 交互,计算成本随输入长度呈 平方级 增长。将上下文长度加倍,可能会让 TTFT 增加四倍。

对于 RAG 来说,这种方式尤其浪费,原因有三:

  1. 信息稀疏: RAG 提示包含用户查询和多个检索到的段落。其中许多段落是冗余的或仅部分相关,但仍然消耗相同的计算时间。
  2. 非结构化注意力: 检索过程中会促进多样化并对文本块去重。同一段落内的 token 之间有很强的注意力,而跨段落的注意力则很弱。这导致生成一种 块对角 注意力模式,其中许多非对角元素的计算几乎无效。
  3. 知识被丢弃: 检索段落时,检索器已经计算了嵌入。通常,这些预计算的信息会被丢弃,迫使 LLM 重新编码文本。

LLaMA-2-7B-Chat 模型中不同检索段落的注意力值可视化。强烈的对角线模式表明注意力主要集中在单个段落内部,这凸显了计算段落间注意力的冗余性。

图 1: 不同检索段落的注意力值。对角线热图模式 (P0–P4) 显示段落内注意力高、段落间注意力低,意味着存在大量计算浪费。

REFRAG 的见解既激进又简单:** 如果对 RAG 上下文的大部分计算都是不必要的,那就去掉它。**


核心思想: 压缩、感知与扩展

REFRAG 改变了传统做法,它不是将原始 token 输入 LLM,而是输入检索段落的 压缩表示。它使用一个轻量级编码器模型,将每个文本块压缩成一个 块嵌入 (chunk embedding)

流程概述:

  1. 分块 (Chunking) : 将检索得到的上下文拆分成固定大小的块 (例如 k=16 个 token) 。
  2. 编码与压缩 (Encoding & Compression) : 使用轻量级编码器 (如 RoBERTa) 将每个块转换为一个嵌入。当 k=16 时,上下文长度缩短 16 倍。
  3. 解码 (Decoding) : 将压缩后的块嵌入映射到解码器的 token 空间,并与用户查询的 token 嵌入一起输入解码器。

REFRAG 的主要架构。轻量级编码器将文本块压缩成嵌入,然后输入主解码器 LLM,从而显著减少输入长度和延迟。

图 2: REFRAG 架构。上下文被分块、编码成可预计算的嵌入,并在解码前由强化学习策略选择性扩展。

压缩大幅缩短了解码器的输入,从而显著降低 TTFT。块嵌入可以预先计算并缓存,实现跨查询复用相同段落的能力。


回报: 在不损失精度的情况下提升速度

基准测试显示,REFRAG 的性能优于标准 LLaMA 和此前的最优方案 CEPE。当 k=16 时,REFRAG 在 16k-token 上下文上 (缓存模式) 可实现 16.53 倍 TTFT 加速,而当 k=32 时,加速比达到 30.85 倍

图表展示了 REFRAG 与基线模型的推理加速对比。REFRAG 的 TTFT 加速随着输入 token 数的增加而显著提升。

图 3: k=16 时的推理加速。REFRAG (蓝色/绿色) 在 TTFT 加速方面远超 CEPE (红色) 。

关键是,这些速度提升并未带来困惑度或任务准确率的显著下降——这得益于专门的训练方法。


训练秘诀: 教会 LLM 理解压缩

直接将压缩嵌入输入未经修改的 LLM 是不行的。REFRAG 的作者设计了一个两阶段对齐策略:

阶段一: 重建任务

  • 冻结解码器。
  • 训练编码器和投影层。
  • 任务: 将一个块编码成嵌入,然后由冻结的解码器根据该嵌入重建原始 token。
  • 目标: 让编码器输出与解码器的预期对齐。

阶段二: 持续预训练的课程学习

  • 解冻解码器。
  • 在压缩块的基础上进行下一段落预测任务,同时训练编码器与解码器。
  • 从简单任务开始 (仅一个块的预测) ,逐步增加到多块、长上下文的预测。

课程学习期间数据混合的可视化。训练过程逐渐从易到难,从少量块重建过渡到更多块重建。

图 4: 课程学习在不同阶段,从短/易的上下文过渡到更长、更复杂的上下文。

这种分阶段策略使解码器能有效解释压缩输入,在高压缩率下仍然保持精度。


更进一步: 基于强化学习的选择性压缩

统一压缩可能丢失关键信息。为解决这一问题,REFRAG 引入了 选择性压缩——保留重要块的原始 token 形式。

一个轻量级的强化学习 (RL) 策略会决定哪些块需要扩展。该策略通过最小化困惑度来最大化预测精度,学会选出对任务最重要的“VIP”块。

选择性 token 压缩示意图,其中 RL 策略决定使用压缩嵌入或完整 token 表示。

图 5: 选择性 token 压缩。RL 策略扩展关键信息块,其余块保持压缩。

不同选择性压缩策略的比较。RL 策略在各种压缩率和数据集下困惑度表现均优于启发式和随机策略。

图 6: 不同数据集与压缩率下,RL 策略在困惑度上优于启发式和随机方法。


实验结果

困惑度与长上下文处理

REFRAG 模型在性能上稳定优于 CEPE、REPLUG 以及截断上下文的 LLaMA 基准模型。它能用一个上下文长度限制为 4,096 的基础 LLaMA,处理长达 16,384 个 token 的上下文——相当于将窗口扩展了 4 倍

长上下文任务中各模型困惑度结果表。REFRAG 模型性能优于其他高效基线,并接近全上下文模型。

表 1: 不同上下文长度下的困惑度结果。REFRAG_8 和 REFRAG_16 优于其他高效基线。


RAG 应用性能

在 16 个不同问答数据集上,REFRAG 也表现优异。在相同延迟下,压缩技术让 REFRAG 可以处理更多段落。例如:

  • LLaMA 可处理 1 个段落
  • REFRAG_8 在相同时间可处理 8 个段落,从而获得更高准确率

REFRAG 和 LLaMA 在性能-延迟、性能-段落数对比中的表现。在相同延迟下,REFRAG 性能显著更优。

图 7: RAG 性能。在段落数相同时,二者相近;但在延迟相同的情况下,REFRAG 可处理更多上下文,从而准确率更高。

在强检索器条件下,REFRAG 在相同延迟下比 LLaMA 平均准确率高 1.22%;弱检索器条件下,这一提升扩大到 1.93%


多轮对话

在多轮 RAG 中,对话会积累上下文历史。标准固定窗口 LLM 必须截断早期轮次,导致丢失重要信息。而 REFRAG 可保留完整历史,从而在多轮中保持准确率。

多轮 RAG 任务上的性能结果表。REFRAG 在对话轮数增加时表现更优。

表 2: 多轮对话结果。REFRAG 避免了截断惩罚,尤其在 ≥6 轮时显著优于微调后的 LLaMA。


结论: 针对特定问题的特定解决方案

REFRAG 证明了 重新思考问题 能够带来巨大收益。通过利用 RAG 稀疏、块对角的注意力结构,REFRAG 将块压缩、专用训练流程以及智能选择性扩展相结合,实现了:

  • 高达 30.85 倍的 TTFT 加速
  • 扩展有效上下文窗口
  • 与全上下文基准相当或更高的准确率

对于延迟敏感、知识密集型的系统——如 RAG 聊天机器人、多轮问答或长文档摘要——REFRAG 展示了你既能获得深度,也能兼顾速度。这是一个有力的例子,说明如何根据数据的真实结构定制效率优化技术,为未来 LLM 推理中的定向优化铺平道路。