Meta最新REFRAG框架引爆RAG圈！KV缓存暴降90%，速度狂飙30×

一、RAG一些棘手的问题

检索增强生成（RAG）通过把外部知识以“拼接文档”形式喂给大模型，在多轮对话、开放域问答、Agent 等场景取得显著效果。然而，该范式带来两个问题：

• 延迟爆炸：输入长度骤增，首词延迟（TTFT）与 KV-Cache 显存随 token 数线性乃至二次方增长。
• 信息稀疏：检索返回的几十篇文档里，仅极少数片段与当前 query 真正相关；其余 token 对生成几乎无贡献，却仍要参与全部注意力计算。

有一些缓解手段，但是效果有限。例如：

• 稀疏注意力 / 线性注意力：把 O(n²) 降为 O(n) 或 O(n log n)，但 KV-Cache 体积不变，TTFT 依旧随序列长度线性增加。
• 前缀压缩 / 滑动窗口：只能丢弃最左侧 token，无法处理“中间大段无用信息”的 RAG 场景。
• Chunk 级编码器（如 CEPE）：把文档先经编码器压成向量，再用交叉注意力注入解码器，但仍以“token”为基本单元，未突破“所有 token 必须进 KV-Cache”的桎梏，且仅支持前缀压缩，不支持多轮、任意位置压缩。

二、4个问题的解决思路

三、REFRAG

3.1 REFRAG的输入与输出

• 输入：question token q₁…qₙ，检索返回文档 d₁…dₘ，总长度 T = n + m。
• 分块：文档侧每 k 个 token 切为一 chunk Cᵢ，得 L = ⌈m/k⌉ 块。
• 编码：轻量编码器 Mₑₙₑ(·)（RoBERTa 系列）输出 chunk embedding cᵢ = Mₑₙₑ(Cᵢ)。
• 投影：φ(cᵢ) → eⁱᶜⁿᵏ ∈ ℝᵈ，与解码器词嵌入同维。
• 解码器输入：[e₁,…,eₙ ⏟ question tokens;  e₁ᶜⁿᵏ,…,eᴸᶜⁿᵏ ⏟ compressed chunks]序列长度由 n + m 缩至 n + L（L ≈ m/k）。
• 生成：标准自回归，注意力只在 n+L 个表示上计算，复杂度 O((n+L)²) ≈ O(n²/k²)。

3.2 训练三阶段（稳定压缩的关键）

1. 重建任务（Reconstruction）目的：让 φ 与 Mₑₙₑ 学会“k 个 token → 一个向量 → 还原 k 个 token”的信息守恒。做法：冻结解码器，仅训练 Mₑₙₑ + φ；输入前 s token，用 chunk 嵌入重建自身。
2. 持续预训练（CPT，Next-Paragraph Prediction）目的：让解码器“习惯”把 chunk 嵌入当上下文。做法：解冻解码器，以“前 s token 的 chunk 嵌入”预测后 o token，采用课程学习：

• 阶段 1 只重建 1 个 chunk，阶段 9 重建 256 个 chunk，逐步加难度，防止梯度爆炸/欠拟合。

3. 下游微调 + 选择性解压（SFT + RL Expansion）目的：适配 RAG、多轮对话、摘要等任务，并学会“重要 chunk 不解压”。做法：

• 先以 90 % 压缩率混合输入（90 % chunk 嵌入 + 10 % 原始 token）继续微调。
• 再用强化学习（PPO-GRPO）训练策略 πθ，按顺序挑选 T′ 个 chunk 恢复成 token；奖励 = -perplexity。推理时，给定延迟预算，可动态选 T′，实现“同延迟更多文档”或“同文档更低延迟”。

3.3 复杂度与加速比

记 s = 文档侧 token 数，k = 压缩粒度，则

• KV-Cache 内存：降为 1/k
• TTFT 计算量：≈ 1/k²（短文本 1/k，长文本 1/k²）
• 实验峰值：k=32 时 TTFT 加速 30.85×，吞吐量最高 6.78×，perplexity 与 Full-Context 基线持平甚至更好。

四、REFRAG的流程

1.检索端改造

• 沿用现有双塔向量检索，仅加“保存 chunk 向量”逻辑；向量库存 FP16，占空间 < 1 % 原文本。

2.编码器离线蒸馏

• 选 RoBERTa-Base 为教师，用 MSE 损失对齐 LLaMA 词嵌入空间；训练 1 epoch，8×A100 约 6 小时。

3.重建任务

• 冻结 LLaMA，仅训编码器+2 层 MLP 投影；学习率 2e-4，batch 256，4 万步即可。

4. CPT + 课程学习

• 数据混合按表 8 几何序列；9 个阶段共 40 k 步，学习率 5e-5，FSDP 8 节点 64 卡，一天完成。

5.下游微调

• RAG 数据 110 万条，混合 5 域问答；90 % 压缩率混合输入，训练 1 epoch，学习率 2e-5。

6.RL 策略训练（可选）

• 用 GRPO 组采样 64 条轨迹，T′=L/4，clip=0.2，训练 5 k 步即可上线；推理时按延迟预算选 T′。

7 部署

• 缓存 chunk 向量放 CPU 内存，推理时异步预取；
• batch=1 实测 A100 上 16 k token 上下文 TTFT 从 2.8 s → 90 ms。

五、实验验证与消融分析

5.1 主实验结果

数据集：Arxiv、Book、PG19、Proof-pile + 16 项 RAG 基准（NQ、TQA、MMLU 等）基线：LLaMA-2-7B、LLaMA-32K、CEPE、REPLUG、LLaMAK（截断匹配 token 数）指标：perplexity、ExactMatch、F1、Accuracy、TTFT、Throughput

结论：在相同或更低困惑度前提下，REFRAG 把首词延迟压到 CEPE 的 1/3.75，显存仅用 3 %。

5.2 下游任务抽样

• RAG（强检索器）：REFRAG8 与 LLaMA-FT 10-passage 相比，平均准确率 +1.22 %，延迟仅 1/5.26。
• RAG（弱检索器）：REFRAG 因能装 8× 更多文档，在检索质量差时优势放大，平均 +1.93 %。
• 多轮对话：TopiOCQA 6 轮场景，REFRAG8 比 LLaMA-FT 提升 5.6 %，因无需截断历史。
• 长文档摘要：Arxiv/Pubmed  Rouge-1 提升 3–5 点，同解码 token 预算下显著领先。

5.3 消融实验

1. 无课程学习：重建任务 perplexity 从 0.135 → 1.599，CPT 不收敛。
2. 无重建预初始化：CPT 阶段 perplexity 差 30 % 以上。
3. 无 RL 选择：固定压缩率 8 的 REFRAG8 被“REFRAG16+RL 选 50 %”全面超越，验证动态解压>静态低压缩。
4. 编码器大小：RoBERTa-Base→Large 仅降 0.2 % perplexity，解码器 7B→13B 降 8 %，说明瓶颈在解码器容量而非编码器。

六、REFRAG与其它方案对比

七、RAG改进的几个方向

1. 极端压缩率（k≥64）下信息损失陡增，需探索非均匀分段或层次压缩。
2. 多模态扩展：文本-图像混合 RAG 时，chunk 定义与对齐策略待研究。
3. 在线学习：随着知识库更新，如何增量更新 chunk 向量而避免全量重算。
4. 端侧落地：chunk 向量缓存仍占内存，可结合量化+索引压缩进一步瘦身。

八、结论

REFRAG 重新审视了 RAG 的“长上下文”假设，指出其注意力矩阵的块对角稀疏性使得“token-level 全计算”成为巨大浪费。

通过“先压缩、再感知、后扩展”，在零解码器参数、零生成质量损失的前提下，把首词延迟压到现有工作的 1/3.75，显存节省 90 % 以上，并支持任意位置、动态压缩的复杂场景。

实验覆盖 RAG、多轮对话、长文摘要三大任务，加速比 6–30× 全面领先，为“大模型+知识库”在 web-scale、Agent 等延迟敏感场景落地提供了可工程化、可扩展的新范式。

​​https://arxiv.org/pdf/2509.01092​​

​​https://github.com/simulanics/REFRAG​​

本文转载自​​​CourseAI​​​，作者：CourseAI