SAGE框架：基于注意力机制的长文档RAG高效压缩方案

SAGE框架注意力机制长文档处理

于 2026-05-31 03:03:13 修改

·本内容遵循CC 4.0 BY-SA版权协议

1. 项目概述：当RAG遇上长文档，我们到底在烦恼什么？

如果你最近在折腾基于大语言模型（LLM）的应用，尤其是检索增强生成（RAG）系统，那你肯定对“长文档”这三个字又爱又恨。爱的是，LLM终于能“看到”更长的上下文了，理论上能处理更复杂的任务；恨的是，把动辄上万token的文档一股脑塞给模型，不仅成本飙升、响应变慢，更糟心的是，答案质量可能不升反降。这就是所谓的“迷失在中间”（Lost in the Middle）现象——模型对输入开头和结尾的信息记得更牢，中间部分则容易忽略。于是，一个核心矛盾摆在我们面前：我们既需要给模型提供足够的上下文来保证答案的准确性和深度，又必须严格控制输入长度以保障效率和成本。

传统的解决方案是RAG：先把长文档切分成块（chunk），用向量数据库检索出与问题最相关的几个块，再把这些块连同问题一起交给LLM生成答案。这个方法简单有效，成了行业标配。但它有个“先天不足”：检索依赖的是向量相似度。如果答案需要综合文档中多个分散的证据（多跳推理），或者问题的表述和文档中的证据在表面词汇上差异很大（语义鸿沟），基于相似度的检索就很容易“翻车”，要么漏掉关键信息，要么塞进来一堆无关的噪音。

最近我在实际项目中就遇到了这种困境。面对几十页的技术白皮书或科研论文，用户问一个需要综合摘要、方法、结果多个部分才能回答的问题时，传统RAG的表现时好时坏，很不稳定。就在我琢磨着是不是要上更复杂的重排序（Re-ranking）或者递归检索（Recursive Retrieval）时，看到了SAGE这个框架。它提出了一种截然不同的思路：不靠向量检索，而是让一个轻量级的“裁判”LLM，通过计算注意力分数，直接告诉我们应该保留文档中的哪些部分。这个想法一下子吸引了我，经过一番研究和实验，我发现它确实在特定场景下，为长文档RAG的效率瓶颈提供了一个非常巧妙的解法。下面，我就结合论文和我的实操经验，来详细拆解一下SAGE到底是怎么工作的，以及我们该如何把它用起来。

2. SAGE核心设计：用注意力机制充当“信息滤网”

SAGE的全称是“基于注意力引导的上下文压缩框架”。这个名字听起来有点学术，但它的核心思想可以用一个比喻来理解：想象你是一位法官，面前堆着如山案卷（长文档），你需要回答一个具体问题（查询）。传统RAG像是一个书记员，他快速翻阅案卷，把所有提到某个关键词的页码都找出来给你。而SAGE则像是一位经验丰富的助理，他会仔细阅读你的问题，然后快速浏览整个案卷，在心里默默评估每一段内容与问题的关联程度，最后只把他认为最相关、最核心的几段话摘要出来，呈到你面前。这位助理做判断的依据，就是“注意力”（Attention）——大语言模型理解文本间关联的核心机制。

2.1 为什么是注意力，而不是向量检索？

要理解SAGE的价值，得先看清传统向量检索的短板。向量检索的本质是“记忆匹配”。它把文本变成高维空间中的点，通过计算点与点之间的距离（如余弦相似度）来判断相关性。这种方法在事实性、片段式的问题上表现很好，比如“文档里某公司的注册日期是什么？”。但是，它难以处理需要“理解”和“连接”的任务。

举个例子，问题可能是：“这篇论文的创新点是什么？” 这个问题的答案可能分散在摘要、引言、方法论和结论等多个部分，并且很少会直接出现“创新点是”这样的字眼。向量检索可能会召回包含“新颖方法”、“主要贡献”等词的段落，但很容易漏掉那些语义相关但用词不同的论述，或者错误地引入一些讨论“相关工作”（即别人的创新）的噪音段落。

注意力机制则不同。当一个小型LLM（我们称之为“注意力模型”）同时阅读问题和文档时，它会动态地计算问题中的每个token与文档中每个token之间的关联强度。这个过程模拟了人类阅读时的聚焦行为：读到问题时，大脑会下意识地在长文中寻找能解答它的线索。这种关联是深层次的、语义上的，而不仅仅是词汇表面的匹配。因此，注意力机制更擅长捕捉那些分散的、需要推理才能关联起来的证据。

2.2 SAGE的工作流程拆解

SAGE的整个流程可以清晰地分为三步，我结合自己的理解画了个简单的示意图在脑子里，这里用文字描述出来：

注意力评分：将用户查询（Query）和整个长文档（Document）拼接在一起，输入到一个轻量级的“注意力模型”（例如LLaMA-3.2-1B）中。这个模型不做生成，只做一件事：计算文档中每个token相对于查询的注意力分数。这个分数直观地反映了“当模型思考这个查询时，它有多关注文档中的这个词”。
分数聚合与排序：文档通常由多个句子或段落组成。SAGE会将属于同一个自然段落或一个固定大小窗口内的token的注意力分数进行聚合（例如取平均或求和），得到这个文本片段的整体相关性分数。然后，所有片段按分数从高到低排序。
预算感知的选择与压缩：用户会设定一个严格的Token预算（比如只允许使用原文10%的token）。SAGE从排名最高的片段开始，依次选取，直到所选片段的token总数达到预算上限。这些被选中的片段，就构成了压缩后的“精华上下文”。最后，将这个压缩后的上下文和原始查询一起，交给另一个更强大的“生成模型”（如Gemini Flash）来生成最终答案。

这个过程最大的优势是精准和可控。精准在于它基于更深层的语义关联进行选择；可控在于它能严格满足用户设定的Token预算，而基于chunk的检索很难精确控制最终prompt的长度。

注意：这里用的“注意力模型”和“生成模型”是分离的。这是一个关键设计。我们可以用一个很小的、专门优化的模型来做注意力计算（成本低、速度快），而把宝贵的计算资源留给最终的答案生成。这种解耦设计让整个系统更加灵活和经济。

2.3 关键技术增强：差分注意力（Differential Attention）

如果SAGE只是简单计算原始注意力，可能会遇到一个问题：有些文档内容本身就很“突出”（比如章节标题、加粗关键词），无论针对什么查询，它们都可能获得较高的注意力分数，这就会引入噪音。

为此，SAGE引入了一个巧妙的技巧：差分注意力。它的思想是，通过引入一个“对比项”来过滤掉那些与查询无关的、但自身很显眼的文本的注意力。

具体实现有两种方式：

固定对比查询：使用一个通用的、无意义的查询（例如“。”），计算文档对该查询的注意力分数。然后用原始查询的注意力分数减去这个“背景”分数，得到差分分数。这相当于去除了文档内容的固有显著性偏差。
最远查询：从一组预设的查询中，选择一个与当前查询语义上“最远”的查询（通过向量相似度判断），用它的注意力分数作为对比。这能更好地抑制那些与当前查询不直接相关，但可能与某些其他查询相关的内容。

在论文的实验中，“最远查询”方法表现最好。这好比在挑选证据时，不仅看它是否与本案有关，还要看它是否与其他无关案件也有关联——如果也有关联，那它作为本案证据的独特性就下降了。

3. 实战解析：如何构建与评估SAGE系统

看懂了原理，接下来我们聊聊怎么把它落地。SAGE论文中使用了多个数据集进行评测，我们正好可以借鉴它的实验设置，来规划自己的实现方案。

3.1 环境与模型选型

首先需要明确几个角色，我建议的选型策略如下：

注意力模型：负责计算注意力分数。核心要求是轻量、高效。论文使用了LLaMA-3.2-1B、Qwen3-8B和Qwen3-14B。对于大多数实际应用，从LLaMA-3.2-1B或Qwen3-1.8B这类小模型开始是完全可行的。它们的计算开销小，且在论文中证明即使只有1B参数，效果也已优于传统RAG。如果你的场景对精度要求极高，且预算充足，可以考虑7B以上的模型。
- 实操心得：注意力模型不一定需要最新的、最强的，但需要支持足够长的上下文窗口（至少要能容纳你的文档+问题）。它的推理速度直接影响整个系统的响应时间，因此需要在效果和效率间权衡。
生成模型：负责根据压缩后的上下文生成答案。核心要求是强大的理解和生成能力。论文统一使用了Gemini-2.5-Flash-Lite以保证对比公平。在实际中，你可以根据任务选择GPT-4、Claude-3、DeepSeek-V2或开源的Qwen-Max等。关键是要固定生成模型，这样才能公平地比较不同上下文选择方法的效果。
嵌入模型：用于“最远查询”差分注意力中的语义距离计算。核心要求是速度快、能较好反映语义相对距离。论文推荐了all-MiniLM-L6-v2，这是一个在MTEB基准上表现良好且非常高效的句子嵌入模型。它不适合做最终的检索，但用于计算查询间的相对相似度绰绰有余。

3.2 分步实现指南

假设我们有一个长文档doc.txt和一个用户查询query，目标是使用SAGE框架来获得答案。

步骤一：文档预处理与分片 虽然SAGE最终选择的是token级片段，但直接对数十万token的全文计算注意力开销巨大。一个实用的策略是进行初步分片。

使用文本分割器（如LangChain的RecursiveCharacterTextSplitter）将长文档按段落或固定大小（如512或1024个字符）分割成较大的“块”。
这一步的目的不是用于检索，而是为了在计算注意力时进行分批处理，避免一次性输入过长。记录每个块在原文中的起止位置。

步骤二：计算注意力分数 这是SAGE的核心步骤，我们需要编写一个函数来计算每个文本片段的注意力分数。

PYTHON

import torch

from transformers import AutoTokenizer, AutoModelForCausalLM

def compute_attention_scores(query, document_chunks, attention_model_name="meta-llama/Llama-3.2-1B"):

"""

计算查询对每个文档块的注意力分数。

简化版，实际需考虑长文档分批、注意力头聚合等细节。

"""

tokenizer = AutoTokenizer.from_pretrained(attention_model_name)

model = AutoModelForCausalLM.from_pretrained(attention_model_name, torch_dtype=torch.float16, device_map="auto")

all_scores = []

for chunk in document_chunks:

# 将查询和文档块拼接

input_text = f"Query: {query}\nDocument: {chunk}"

inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=4096).to(model.device)

with torch.no_grad():

outputs = model(**inputs, output_attentions=True)

# 获取最后一层注意力权重，形状为 (batch, heads, seq_len, seq_len)

attentions = outputs.attentions[-1]

# 我们关心查询token对文档token的注意力。简化处理：取平均

# 假设输入序列中，查询部分在前，文档部分在后。

query_len = len(tokenizer.encode(f"Query: {query}"))

chunk_attention = attentions[:, :, :query_len, query_len:].mean(dim=[0,1,2]) # 平均批次、头、查询维度

chunk_score = chunk_attention.mean().item() # 得到该块的整体注意力分数

all_scores.append(chunk_score)

return all_scores

重要提示：上面的代码是一个高度简化的示意。实际实现中，你需要精确处理token位置，正确提取查询到文档的注意力权重，并可能需要对多个注意力头的结果进行聚合（例如取最大值或L2范数）。此外，对于超长文档，需要更精细的流式处理来避免内存溢出。

步骤三：差分注意力计算（可选但推荐） 为了提高选择质量，实现差分注意力。

准备一个对比查询列表，可以包含一些通用查询如“。”, “Summarize.”, “What is this document about?”。
使用嵌入模型（如all-MiniLM-L6-v2）计算当前查询与列表中每个查询的相似度，选择相似度最低的作为“最远查询”。
用同样的compute_attention_scores函数计算文档对“最远查询”的注意力分数。
将原始查询的注意力分数减去最远查询的注意力分数（或进行其他标准化处理），得到差分分数作为最终的相关性分数。

步骤四：预算感知的选择 根据预算选择片段。

将文档块（或更细粒度的句子）按照其最终的注意力分数（或差分分数）降序排列。
从分数最高的块开始，依次将其加入候选列表，并累加其token数量。
当累加的token数达到用户预设的预算（如文档总token数的10%）时停止。
将被选中的块按照它们在原文中的顺序拼接，形成压缩后的上下文。

步骤五：答案生成 将压缩后的上下文和原始查询，按照预设的提示词模板格式化，发送给生成模型获取最终答案。

PYTHON

prompt_template = """请基于以下上下文回答问题。如果上下文不包含相关信息，请回答“根据提供的信息无法回答”。

上下文：

{compressed_context}

问题：{query}

答案："""

final_prompt = prompt_template.format(compressed_context=compressed_context, query=query)

# 调用生成模型API或本地模型

answer = generate_with_llm(final_prompt, model="gemini-2.0-flash")

3.3 效果评估：不只是看准确率

论文从多个维度评估SAGE，我们在实际项目中也可以借鉴这套方法，不要只看一个最终答案的对错。

LLM即裁判（LLM-as-a-Judge）：对于开放域问答，答案可能有多样性。可以请一个强大的LLM（如GPT-4）作为裁判，根据问题和参考答案，对SAGE和基线方法生成的答案进行二分类判断（正确/错误）。这比精确匹配（Exact Match）更合理。
语义相似度：使用嵌入模型（如Qwen3-Embedding-0.6B）计算生成答案与标准答案的向量余弦相似度。这个指标能反映答案在语义上的贴近程度，尤其适合评估开放性、论述性问题的回答质量。
Token效率：这是SAGE的核心优势。在相同的答案质量下，比较SAGE和基线RAG所消耗的上下文Token数量。通常可以绘制“准确率/相似度 vs. Token使用量”曲线，SAGE的曲线应该更靠左上角（即用更少的Token达到相同或更高的质量）。
推理速度：分别测算注意力计算、检索（如果是RAG）、生成等环节的耗时。对于SAGE，注意力模型的前向传播是主要开销，需要评估其在实际硬件上的延迟。

4. 多场景实测：SAGE的优势与边界在哪里？

论文在四个数据集上进行了测试，这几乎覆盖了长文档QA的主要场景。我们来逐一解读，这能帮我们判断SAGE适合用在哪里。

4.1 场景一：复杂推理阅读理解（QuALITY-hard）

数据集特点：文章长，问题需要多步推理，答案可能分散在全文。
核心发现：SAGE在此场景下优势最大。在仅使用10% Token预算的情况下，SAGE（使用Qwen3-14B注意力模型）能达到与使用近40% Token预算的强RAG基线相近的准确率。即使只用1B参数的注意力模型，SAGE也全面超越了所有RAG基线。
我们的启示：当你的任务涉及复杂的逻辑推理、信息整合，而不仅仅是事实查找时，SAGE是比传统RAG更优的选择。 因为它通过注意力机制捕捉到了证据之间的语义关联链。

4.2 场景二：科学论文理解（Paper）

数据集特点：文档极长，结构清晰（摘要、引言、方法等），问题类型多样（从事实型“发表年份”到语义型“主要贡献”）。
核心发现：SAGE同样有效，且无需针对论文结构进行任何调整。对于“主要贡献”这类开放性问题，SAGE的性能随Token预算增加而稳定提升，而RAG提升缓慢。这说明SAGE选取的额外Token确实是“有帮助的”，而不是噪音。
我们的启示：对于技术文档、研究报告、法律条文等具有内在逻辑结构的长文档，SAGE的注意力机制能很好地理解其语义组织，从而精准定位答案所在的部分。 差分注意力技术在这里尤其有用，可以过滤掉像“参考文献”、“附录”这类与具体问题无关但格式突出的部分。

4.3 场景三：法规通知解析（Notice）

数据集特点：文档长度相对均匀，便于进行严格的Token预算对齐对比。
核心发现：SAGE能在低于单个RAG分块（chunk）大小的极端预算下工作（例如5%的Token），而RAG由于最小检索单元是一个chunk，无法支持如此精细的预算控制。这证明了SAGE细粒度选择的能力。
一个有趣的失败案例：在查询“公司所在州缩写和邮编”时，SAGE在预算很小时表现最好，但随着预算增加，准确率反而下降，逐渐被RAG追平。原因是文档中可能存在多个邮编，当预算增加，SAGE可能把其他无关的邮编也选了进来，造成了混淆。
我们的启示：SAGE在极低预算场景下具有不可替代的优势。但同时，它也暴露了一个弱点：当文档中存在大量相似或重复实体时，贪婪的注意力选择策略可能会引入歧义。这提示我们，对于此类文档，可能需要结合实体识别或后处理去重来优化。

4.4 场景四：表格问答（AIT-QA）

数据集特点：数据是结构化的JSON表格，传统RAG的文本分块会破坏表格结构。
核心发现：SAGE可以自然地扩展到结构化数据。研究者将计算单位从文本片段改为表格的“行”，计算查询对每一行数据的注意力，然后选择最相关的行。实验表明，仅选择最相关的4行（占用约51%的Token），就能达到与使用整个表格（100% Token）相近的准确率（0.87 vs 0.88）。
我们的启示：SAGE的框架是通用的，其核心是“基于注意力的重要性评估”。 只要能为你的数据（文本、表格、甚至未来可能的多模态数据）定义出合理的“可选择性单元”（句子、段落、行、区域），并能够计算注意力，就可以应用SAGE的思想进行压缩。

5. 避坑指南与进阶思考

在实际尝试复现和运用SAGE的过程中，我踩过一些坑，也产生了一些延伸思考。

5.1 常见问题与排查

注意力分数全都很低或没有区分度
- 可能原因：注意力模型太小或与任务不匹配；查询和文档的拼接格式不符合模型的训练方式；没有正确提取查询到文档的注意力权重。
- 排查步骤：首先检查输入格式，确保其类似于模型预训练时的常见格式（如[INST] 问题 [/INST] 文档）。其次，手动检查几个片段，看看模型是否真的在“理解”内容。可以尝试换一个稍大一点的注意力模型（如从1B到3B）做对比实验。
差分注意力效果不明显甚至变差
- 可能原因：对比查询选择不当。“固定对比查询”如果太通用，可能无法有效过滤特定噪音。“最远查询”如果选到的查询与当前查询在某个维度上意外相关，则会产生误导。
- 解决方案：可以尝试构建一个更丰富、更多样化的对比查询池。或者，暂时放弃差分注意力，先使用原始注意力，观察效果。有时，在文档噪音不大的情况下，原始注意力已经足够好。
生成答案时出现“幻觉”，引用未被选中的内容
- 可能原因：压缩过程过于激进，丢失了必要的背景信息或指代关系。例如，压缩后的段落里出现“上述方法”，但“上述方法”的具体描述在前文已被裁剪。
- 解决方案：适当增加Token预算。或者在选择片段时，不仅考虑片段本身，还考虑其与已选片段的连贯性，可以尝试在排序后，对选中的片段进行轻微的上下文扩展（如前缀后加几个句子）。
处理超长文档时速度慢
- 可能原因：注意力模型需要处理整个文档，即使分批，计算量也很大。
- 优化思路：采用两阶段策略。第一阶段，用一个非常快速的方法（如基于BM25的关键词检索或一个小型嵌入模型）粗筛出可能相关的较大区域。第二阶段，只在这个缩小的区域上运行SAGE进行精筛。这本质上是“检索+SAGE”的混合架构。

5.2 成本与延迟的权衡

SAGE引入了一个额外的注意力模型前向传播开销。论文中的数据显示，使用LLaMA-3.2-1B处理整个文档，其速度与使用某些中型嵌入模型（如Qwen3-Embedding-8B）进行全量向量检索处于同一量级，但比轻量级嵌入模型（如UAE-Large-V1）要慢。

决策点：如果你的应用对延迟极其敏感（<100ms），且问题多为简单事实型，传统RAG搭配轻量嵌入模型可能仍是首选。如果你的应用可以接受稍高的延迟（几百毫秒到秒级），且问题复杂、文档专业，那么SAGE带来的答案质量提升将是值得的。关键在于评估“质量提升”带来的业务价值是否大于“延迟增加”带来的体验成本。

5.3 与现有RAG系统的融合

SAGE并非要完全取代RAG，它可以作为RAG系统中的一个高级“精排”或“重压缩”模块。

方案A（先检索后压缩）：先用快速向量检索召回Top-K个相关chunk。然后将这K个chunk拼接成一个“中等长度”的上下文，再用SAGE对这个上下文进行二次压缩，得到最终送入生成模型的精华内容。这既利用了检索的快速初筛，又发挥了SAGE的精准选择能力。
方案B（作为重排序器）：用检索召回多个chunk，但不用它们直接拼接。而是将每个chunk单独与查询组合，用SAGE的注意力模型计算每个chunk的注意力分数，然后按此分数对chunk进行重排序，选择Top-N。这种方式比方案A的计算量更大，但可能更精准。

我个人在项目中尝试过方案A，在处理百页级技术手册时，效果比纯检索或纯SAGE都要稳定。先用检索确保不遗漏大范围的相关区域，再用SAGE做精细过滤，在可控的延迟内达到了最佳的质量。

SAGE框架为我们打开了一扇新窗：不再仅仅依赖向量空间的几何距离，而是利用LLM自身的注意力机制来理解相关性。它尤其适合那些答案深藏在冗长、复杂文档中的场景。当然，没有银弹，它带来了额外的计算复杂度和对注意力模型设计的考量。但在长文档信息提取这场攻坚战中，SAGE无疑提供了一把锋利的新武器。我的体会是，在构建生产级系统时，不必拘泥于一种技术，将SAGE与传统的检索技术有机结合，往往能取得“1+1>2”的效果。下一步，我打算探索如何将这种注意力引导的机制，应用于多轮对话中对历史上下文的动态压缩，这或许是另一个值得深挖的方向。