大型语言模型 (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 来说,这种方式尤其浪费,原因有三:
- 信息稀疏: RAG 提示包含用户查询和多个检索到的段落。其中许多段落是冗余的或仅部分相关,但仍然消耗相同的计算时间。
- 非结构化注意力: 检索过程中会促进多样化并对文本块去重。同一段落内的 token 之间有很强的注意力,而跨段落的注意力则很弱。这导致生成一种 块对角 注意力模式,其中许多非对角元素的计算几乎无效。
- 知识被丢弃: 检索段落时,检索器已经计算了嵌入。通常,这些预计算的信息会被丢弃,迫使 LLM 重新编码文本。
图 1: 不同检索段落的注意力值。对角线热图模式 (P0–P4) 显示段落内注意力高、段落间注意力低,意味着存在大量计算浪费。
REFRAG 的见解既激进又简单:** 如果对 RAG 上下文的大部分计算都是不必要的,那就去掉它。**
核心思想: 压缩、感知与扩展
REFRAG 改变了传统做法,它不是将原始 token 输入 LLM,而是输入检索段落的 压缩表示。它使用一个轻量级编码器模型,将每个文本块压缩成一个 块嵌入 (chunk embedding) 。
流程概述:
- 分块 (Chunking) : 将检索得到的上下文拆分成固定大小的块 (例如
k=16
个 token) 。 - 编码与压缩 (Encoding & Compression) : 使用轻量级编码器 (如 RoBERTa) 将每个块转换为一个嵌入。当
k=16
时,上下文长度缩短 16 倍。 - 解码 (Decoding) : 将压缩后的块嵌入映射到解码器的 token 空间,并与用户查询的 token 嵌入一起输入解码器。
图 2: REFRAG 架构。上下文被分块、编码成可预计算的嵌入,并在解码前由强化学习策略选择性扩展。
压缩大幅缩短了解码器的输入,从而显著降低 TTFT。块嵌入可以预先计算并缓存,实现跨查询复用相同段落的能力。
回报: 在不损失精度的情况下提升速度
基准测试显示,REFRAG 的性能优于标准 LLaMA 和此前的最优方案 CEPE。当 k=16
时,REFRAG 在 16k-token 上下文上 (缓存模式) 可实现 16.53 倍 TTFT 加速,而当 k=32
时,加速比达到 30.85 倍。
图 3:
k=16
时的推理加速。REFRAG (蓝色/绿色) 在 TTFT 加速方面远超 CEPE (红色) 。
关键是,这些速度提升并未带来困惑度或任务准确率的显著下降——这得益于专门的训练方法。
训练秘诀: 教会 LLM 理解压缩
直接将压缩嵌入输入未经修改的 LLM 是不行的。REFRAG 的作者设计了一个两阶段对齐策略:
阶段一: 重建任务
- 冻结解码器。
- 仅 训练编码器和投影层。
- 任务: 将一个块编码成嵌入,然后由冻结的解码器根据该嵌入重建原始 token。
- 目标: 让编码器输出与解码器的预期对齐。
阶段二: 持续预训练的课程学习
- 解冻解码器。
- 在压缩块的基础上进行下一段落预测任务,同时训练编码器与解码器。
- 从简单任务开始 (仅一个块的预测) ,逐步增加到多块、长上下文的预测。
图 4: 课程学习在不同阶段,从短/易的上下文过渡到更长、更复杂的上下文。
这种分阶段策略使解码器能有效解释压缩输入,在高压缩率下仍然保持精度。
更进一步: 基于强化学习的选择性压缩
统一压缩可能丢失关键信息。为解决这一问题,REFRAG 引入了 选择性压缩——保留重要块的原始 token 形式。
一个轻量级的强化学习 (RL) 策略会决定哪些块需要扩展。该策略通过最小化困惑度来最大化预测精度,学会选出对任务最重要的“VIP”块。
图 5: 选择性 token 压缩。RL 策略扩展关键信息块,其余块保持压缩。
图 6: 不同数据集与压缩率下,RL 策略在困惑度上优于启发式和随机方法。
实验结果
困惑度与长上下文处理
REFRAG 模型在性能上稳定优于 CEPE、REPLUG 以及截断上下文的 LLaMA 基准模型。它能用一个上下文长度限制为 4,096 的基础 LLaMA,处理长达 16,384 个 token 的上下文——相当于将窗口扩展了 4 倍。
表 1: 不同上下文长度下的困惑度结果。REFRAG_8 和 REFRAG_16 优于其他高效基线。
RAG 应用性能
在 16 个不同问答数据集上,REFRAG 也表现优异。在相同延迟下,压缩技术让 REFRAG 可以处理更多段落。例如:
- LLaMA 可处理 1 个段落
REFRAG_8
在相同时间可处理 8 个段落,从而获得更高准确率
图 7: RAG 性能。在段落数相同时,二者相近;但在延迟相同的情况下,REFRAG 可处理更多上下文,从而准确率更高。
在强检索器条件下,REFRAG 在相同延迟下比 LLaMA 平均准确率高 1.22%;弱检索器条件下,这一提升扩大到 1.93%。
多轮对话
在多轮 RAG 中,对话会积累上下文历史。标准固定窗口 LLM 必须截断早期轮次,导致丢失重要信息。而 REFRAG 可保留完整历史,从而在多轮中保持准确率。
表 2: 多轮对话结果。REFRAG 避免了截断惩罚,尤其在 ≥6 轮时显著优于微调后的 LLaMA。
结论: 针对特定问题的特定解决方案
REFRAG 证明了 重新思考问题 能够带来巨大收益。通过利用 RAG 稀疏、块对角的注意力结构,REFRAG 将块压缩、专用训练流程以及智能选择性扩展相结合,实现了:
- 高达 30.85 倍的 TTFT 加速
- 扩展有效上下文窗口
- 与全上下文基准相当或更高的准确率
对于延迟敏感、知识密集型的系统——如 RAG 聊天机器人、多轮问答或长文档摘要——REFRAG 展示了你既能获得深度,也能兼顾速度。这是一个有力的例子,说明如何根据数据的真实结构定制效率优化技术,为未来 LLM 推理中的定向优化铺平道路。