REFRAG的技术背景
大型语言模型(LLM)在检索增强生成(RAG)任务中表现出色,但其性能严重依赖于输入上下文的长短。然而,处理长上下文会带来两大挑战:
- 时间成本:注意力机制的计算复杂度随序列长度呈平方级增长,导致首Token延迟(TTFT) 极高。
- 空间成本:需要缓存巨大的键值对(KV Cache),内存需求随序列长度线性增长,限制了批处理大小和吞吐量。
现有的长上下文优化方法(如稀疏注意力、流式处理)多为通用方案,未针对RAG场景的独特结构进行优化。而RAG在实际应用中,检索到的多个段落之间往往语义相关性较低,在注意力机制的模型生成token时会呈现“块对角”的稀疏特性(Block-Diagonal Sparsity Pattern),即一个段落内部的token之间会有很强的关联,但不同段落之间的token关联度非常弱。这意味着,将大量原始Token全部输入LLM进行计算是不必要的且低效的。
因此,Meta Superintelligence Labs提出了REFRAG(REpresentation For RAG)技术:通过压缩、感知与选择性扩展上下文,从根本上解决了长上下文输入导致的LLM推理延迟高、内存消耗大的瓶颈问题。REFRAG的论文数据表明,在不损失生成质量的前提下,实现了高达30.85倍的首Token延迟(TTFT)加速,为高性能、低延迟的RAG应用部署提供了切实可行的解决方案。
REFRAG核心思想与架构
REFRAG的核心思想是:用轻量级Encoder模型预先压缩冗长的检索上下文,生成简洁的“块Embeddings”表示,再让LLM主要基于这些压缩表示进行生成,从而大幅减少需要处理的序列长度。
其架构主要由三个核心组件构成:
- 轻量级编码器(Encoder) :一个独立的预训练模型(如RoBERTa),负责将文本块压缩为单个embedding向量。它与主LLM分离,可并行高效计算。
- 投影层(Projection Layer) :一个简单的可训练神经网络(如MLP),用于将编码器Encoder输出的embedding向量映射到主LLM的词向量空间,实现两个模型间的“语言”对齐。
- 主LLM解码器(Decoder) :一个标准的Decoder-only模型(如LLaMA),其本身结构无需任何修改。它的角色从“既编码又解码”转变为主要专注于“解码” ,即根据压缩后的上下文信息生成答案。
REFRAG工程化的详细流程
先训练三个模型
REFRAG在工程化落地时,需要先训练好三个模型或者说三个阶段的模型:
阶段1:重建任务(Reconstruction Task) - “学会压缩”
这个阶段核心目标是训练编码器和投影层,使它们能够以最小化信息损失的方式有效地压缩文本。为实现这一目标,该阶段采用了一种冻结策略,即保持主LLM的参数不变,仅对编码器和投影层进行训练。模型通过尝试从压缩后的块嵌入中重建原始Token来学习有效的表示方法。这一过程的关键在于采用了课程学习策略,模型从重建单个文本块开始,随着训练的推进,逐步增加序列长度和复杂度,最终能够熟练地重建包含多个块的长序列。这种循序渐进的训练方式是模型成功实现有效压缩的重要保证。
阶段2:持续预训练(CPT) - “学会利用压缩信息生成”
在完成压缩能力的学习后,模型进入持续预训练阶段。此阶段的目标是让整个系统学会如何利用压缩后的上下文信息进行自然语言生成。与第一阶段不同,此时解冻了LLM的参数,使其与编码器、投影层进行联合训练。训练任务采用标准的下一段落预测目标,使LLM逐渐学会理解和处理"块嵌入"这种新型的输入形式,从而能够基于压缩的上下文信息生成连贯且准确的文本。
阶段3:强化学习策略(RL Policy) - “学会选择”
第三阶段引入了强化学习策略,旨在训练一个智能的选择机制。该阶段通过一个轻量级Transformer网络作为策略网络,以所有块嵌入作为输入,采用PPO算法进行训练。其奖励函数基于生成结果的负困惑度设计,激励策略网络学会识别哪些信息块对生成质量至关重要,从而决定需要将哪些块展开为原始Token。这个机制使模型能够在保持高效的同时,不丢失关键信息。
REFRAG的推理流程
1 查询与检索
推理流程始于用户提交的一个自然语言问题 Q 。系统随即调用外部的检索系统,从知识库中快速查找并返回与问题 Q 最相关的 K 个文本段落。本步骤的输出结果是一个由这 K个段落拼接而成、包含大量 Token 的长上下文 C ,它将作为后续处理的输入材料。
2 分块与并行编码
接下来,系统对获得的长上下文 C 进行处理。首先,将其分割成 L 个长度固定为 k 的文本块。随后,每个文本块被并行地送入已训练好的轻量级编码器(如 RoBERTa)。编码器独立处理每个块,该向量紧接着通过已训练好的投影层(MLP),被映射到主 LLM 的词嵌入空间,生成最终的“块嵌入” 。此步骤的核心优势在于其高度并行化的处理能力,且生成的块嵌入可以预先计算并缓存,从而极大提升了整体推理速度。
3 智能选择决策
在此步骤中,所有块的嵌入向量被一同输入已训练好的强化学习(RL)策略网络。这个轻量级的网络会评估每个块的重要性,并输出一个决策序列,精准指示哪些块包含至关重要、不容压缩的细节信息(如精确的数字、日期、公式或核心论点等)。
4 构建混合输入序列
这是实现高效推理的核心步骤。系统将按严格顺序拼接三种元素,以构建一个全新的、极短的输入序列:
问题Token嵌入:用户原始问题 Q 的 Token 经过 LLM 本身嵌入层得到的向量。被选中块的原始Token:对于策略网络选出的关键块,将其完整展开,使用其原始 Token 的嵌入向量。被忽略块的压缩嵌入:对于其余的非关键块,则使用其压缩后的块嵌入作为代表。
最终,这样一番操作之后,原本极其冗长的输入文本就被 缩短了,变得非常精练。
5 高效生成答案
最后,将步骤四构建好的混合序列输入已训练好的主 LLM 解码器(如 LLaMA)。LLM 以自回归的方式处理这个短序列。关键在于,由于 LLM 在持续预训练(CPT)阶段已经过充分训练,它完全理解“块嵌入”的语义。它能从压缩嵌入中提取段落级概要,并从展开的原始 Token 中捕捉细微信息。最终,模型快速生成一个流畅且准确的答案 A,并返回给用户,完成整个高效推理流程。
REFRAG优缺点
优点
缺点
| 缺点 | 具体表现与数据指标 | 潜在的影响与成本 |
|---|---|---|
| 复杂的训练流程 | 需要连续进行3个不同阶段的训练(重建、CPT、RL),整个流程耗时漫长。 | 极大地增加了前期研发和时间成本,对技术团队的专业能力要求高,试错成本大。 |
| 高昂的训练成本 | 训练涉及同时优化LLaMA-7B、RoBERTa-Large等大模型参数,计算资源消耗远超单一模型微调。 | 需要大量的GPU算力支持,总训练成本($)非常高昂,提高了技术使用的门槛。 |
| 系统复杂度高 | 系统需维护3个以上的主要组件(编码器、投影层、RL策略网、LLM),并管理其之间的数据流。 | 增加了系统维护、更新和调试的复杂性,引入了更多的潜在故障点,运维负担加重。 |
| 额外的存储开销 | 为加速推理,需为知识库中数百万甚至数亿个文本块预计算并存储其嵌入向量。 | 需要巨大的磁盘存储空间来容纳这些缓存嵌入,增加了硬件和管理的间接成本。 |
| 编码器性能瓶颈 | 编码器的压缩是有损的,其质量上限决定了信息保留的上限。若编码器性能不佳,信息会永久丢失。 | RL策略无法恢复编码阶段已丢失的细节,可能导致关键信息遗漏,影响答案质量。 |
| 依赖上游质量 | 系统表现高度依赖于检索器返回段落的相关性。若检索质量差,即使REFRAG本身高效,最终答案也会错误。 | 错误会沿流程传播和放大,REFRAG的高效反而可能更快地产生错误结果。 |
| 场景局限性 | 对于上下文很短(<1k tokens) 或需要逐字符精确匹配的任务,其优势无法体现。 | 在这些场景下,REFRAG的系统复杂度反而会成为负担,不如直接使用原始模型简单可靠。 |
小结
REFRAG作为Meta超级智能实验室的首篇论文,确实展现出了卓越的工程创新性,通过引入外部轻量编码器对检索上下文进行并行压缩与智能选择,将长文本处理延迟降低了一个数量级,同时在多项任务上保持无损性能,为大规模、低延迟RAG应用落地提供了有价值的新的解决方案。
不过工程化落地还是有诸多困难,虽然“Rethinking”了,但是没开源,短时间也不好验证效果好坏。