Menu
Article RAG

RAG 进阶之路7：REFRAG 机制应对长上下文挑战

Gloary Lei

• 7 min read

之前的文章中，我们从基础设施（Milvus）、语义核心（BGE）、到代码落地（LangChain）以及应用层优化（Small-to-Big/HyDE），构建了一套完整的 RAG 知识体系。

然而，AI 领域的变化是指数级的。随着 Claude 3 支持 200k 上下文，Gemini 1.5 Pro 甚至支持到 1M token，一种论调开始流行：“RAG 已死，Long Context (长上下文) 才是未来。”

毕竟，如果能把整本《红楼梦》或整个公司的知识库直接塞进 Prompt 里，还需要费劲地做切片、建索引、搞检索吗？

但现实是骨感的。“能放进去”不代表“能跑得动”。
把海量文本直接喂给 LLM 会带来两个致命问题：

  1. 首字延迟 (TTFT) 爆炸： Transformer 的计算复杂度是O(N2)。输入长度增加 10 倍，推理时间可能增加 100 倍。
  2. 显存 (KV Cache) 枯竭： 存储长文本的中间状态（Key-Value Cache）需要消耗海量 GPU 显存。

因此，RAG 不会死，它正在进化。下一代的 RAG 不再仅仅是简单的“搜索+拼接”，而是深入模型底层，追求极致的压缩与效率。今天我们要解析的 REFRAG ，就是这种进化方向的代表。

一、 核心背景：为什么需要 REFRAG？

在传统的 RAG 中，检索到的 Top-K 文档（Chunks）会被直接拼接成 Prompt 喂给大模型。

  • 问题 1：计算复杂度爆炸。 如果检索了 10 个文档，每个 500 Token，总长 5000 Token。计算量非常大。
  • 问题 2：信息稀疏。 文档 A 通常只和文档 A 内部的内容相关，和文档 B 没啥关系。 让模型计算文档 A 和文档 B 之间的注意力往往是浪费算力。

二、 REFRAG 的工作机制详解

REFRAG 的核心思想可以概括为：“先看缩略图，看不清再看原图”。

它不把所有原始 Token 塞进模型，而是使用一种混合表示策略。

1. 三步走流程：Compress-Sense-Expand

  • Compress (压缩 - 制作缩略图):
    • 对于检索到的每一个文本块（Chunk），预先计算一个压缩嵌入（Compressed Chunk Embedding）。
    • 这个嵌入就像是一个“软 Token”（Soft Token），它用一个向量代表了原本几百个 Token 的语义概括。
    • 效果： 输入长度从几千个 Token 瞬间缩减为几十个 Embedding 向量。
  • Sense (感知 - 决定是否够用):
    • 这里引入了一个强化学习（RL）代理（Policy）。
    • 它会“看”一眼当前的压缩嵌入和 Query，判断：仅凭这个压缩向量，足够回答问题吗？或者足够预测下一段内容吗？
  • Expand (扩展 - 还原细节):
    • 如果 RL 觉得不够（Information deficient）： 将该块还原为原始的完整 Token 序列（Full Tokens）。
    • 如果 RL 觉得够了（Sufficient）： 直接将压缩嵌入喂给 LLM 的解码器，跳过原始 Token 的计算。

2. 性能提升原理

  • KV Cache 节省： 显存占用大幅减少，因为不需要存储所有原始 Token 的 Key-Value 对，只需存储压缩向量的。
  • 计算量降低：因为 Block 的数量远小于 Token 的数量，这在长上下文场景下是数量级的提升。

3. 实现方案

  1. 基础模型 (Base LLM):
    • 选择 Llama-3-8B, Qwen2-7B。这些模型效果好且显存占用可控。
    • 框架：使用 PyTorch 和 Hugging Face Transformers 库。
  2. 实现步骤：
    • Step 1: 训练压缩器 (Auto-Encoder):
      你需要训练一个小的编码器，把一段文本（比如 512 tokens）压缩成几个 dense vectors（比如 4 个 vectors）。这通常需要大量文本数据进行自监督训练。
    • Step 2: 修改 Attention Mask:
      在 Hugging Face 的模型代码中，重写 forward 函数。利用“块对角线”原理，修改 Attention Mask，让模型只能看到被允许看到的块。
    • Step 3: 插入 RL Policy:
      在模型生成每一个 token 之前，插入一个极小的分类器（Policy Network），判断当前是否需要解压那个 Vector。
    • Step 4: 软提示注入 (Soft Prompt Injection):
      将压缩后的 Vector 直接拼接到 Embedding 层，而不是经过 Tokenizer。

三、 关键难点与风险控制：精度 vs 速度

对精度极其敏感的场景存在的风险

1. 风险点：压缩带来的模糊性

  • 数字敏感度丢失： 比如“税率为 5%”。压缩后的向量可能只保留了“关于税率”的语义，但丢掉了“5%”这个精确数字。
  • 法律条款歧义： 法律条文中“应”和“可”一字之差谬以千里。压缩表示很难保留这种细微的字面差异。

2. 解决方案：RL 策略的训练目标

REFRAG 的成败取决于那个 RL 策略（Gatekeeper）聪不聪明。

  • 负奖励机制（Negative Reward）： 训练时使用“下一段落预测困惑度（Perplexity）”作为惩罚。
    • 如果 RL 决定使用压缩向量，结果导致模型预测下一个词时非常困惑（预测不准），RL 就会收到惩罚。
    • 结果： RL 学会了在遇到包含具体数字、专有名词、复杂逻辑的段落时，倾向于选择 Expand（展开原始 Token）；而在遇到背景介绍、通用描述时，选择 Compress（压缩向量）。

四、 横向对比：REFRAG 在 RAG 优化谱系中的位置

为了让你更全面地理解，我们将 REFRAG 放在 RAG 优化的全景图中：

1. 上下文压缩/选择类技术 (Context Compression/Selection)

REFRAG 属于这一类，主要竞争技术包括：

  • LLMLingua / LongLLMLingua:
    • 原理： 计算每个 Token 的自身信息熵（perplexity），直接删除那些“废话”Token（如 the, a, is），保留关键词。
    • 对比 REFRAG： LLMLingua 是离散的（删 Token），REFRAG 是连续的（向量化压缩）。REFRAG 的理论上限更高，但实现难度更大。
  • Reranking (重排序):
    • 原理： 使用 Cross-Encoder 对 Top-K 结果精排，只选 Top-3 给大模型。
    • 对比 REFRAG： Reranking 是粗粒度的“要或不要”，REFRAG 提供了中间态“要，但只要摘要”。

2. 索引优化 (Pre-Retrieval)

  • Small-to-Big (Parent Document Retrieval): 索引时切小块，检索到小块后，把它的“父文档”（大块上下文）喂给 LLM。
  • HyDE (假设性文档嵌入): 生成一个虚构的完美答案，用这个答案去检索。

五、 总结与实施建议

REFRAG 是 RAG 追求极致性能的高级手段。

  • 优点： 能在保持较高精度的前提下，显著降低首字延迟（TTFT）和推理成本。
  • 缺点： 架构复杂，需要训练额外的 RL 策略模型；需要预先计算并存储 Chunk Embeddings（增加了存储成本）。

对于精度敏感的场景，建议如下：

  1. 基线方案： 先做好 Hybrid Search (关键词+向量) + Reranking (重排序)。这是保证查准率的基础。
  2. 进阶优化： 如果上下文依然过长（例如需要同时对比 10 份法律文件），再考虑引入类似 REFRAG 的上下文压缩技术。
  3. 安全网： 在涉及数字（金额、日期、比例）的查询中，可以通过 Prompt Engineering 强制要求模型引用原始文本，或者强制 RL 策略在检测到 Query 包含“多少”、“时间”等词时，强制展开所有 Block，不进行压缩，以牺牲速度换取绝对精度。
Share