1. 项目概述当长文档问答遇上“信息过载”处理一份动辄几十页、上百页的PDF报告或者一本电子书然后向AI模型提出一个具体问题这可能是很多研究者和开发者正在尝试的事情。理想很丰满我们把整个文档“喂”给模型它就能像一位通读了全文的专家给出精准的答案。但现实很骨感当你真的把一篇数万token的文档塞进提示词Prompt时往往会发现模型的表现不尽如人意——它可能答非所问遗漏关键细节或者干脆因为上下文长度限制而拒绝处理。核心矛盾就在于我们提供给模型的“上下文窗口”是有限的而长文档的信息是冗余且分散的。这就是“SAGE基于注意力机制的上下文压缩框架”要解决的核心问题。它不是一个全新的模型而是一个精巧的“预处理”框架。你可以把它想象成一位高效的“研究助理”在将长文档交给大语言模型LLM这位“主分析师”之前这位助理会先快速浏览全文然后不是原封不动地递交所有材料而是整理出一份高度凝练、只包含与当前问题最相关信息的“摘要简报”。SAGE的创新之处在于它借鉴了Transformer模型核心的“注意力机制”思想来模拟这个“快速浏览并抓取重点”的过程从而实现自适应的、动态的上下文压缩。对于需要处理长文本问答Long-Document QA的开发者、知识库构建者或是任何受限于模型上下文长度但又要保证回答质量的应用场景理解SAGE的原理和实现思路远比简单地调用一个API更有价值。它能帮你从根本上优化信息流用更低的计算和成本开销换取更精准、更可靠的问答效果。2. SAGE框架的核心设计思想与工作原理2.1 问题根源为什么直接输入长文档会失效要理解SAGE的价值首先得明白为什么“暴力”输入长文档效果不好。这背后有几个关键原因计算复杂度与成本Transformer模型的自注意力机制的计算复杂度与序列长度的平方成正比。这意味着将上下文长度从1K扩展到8K计算负担可能增加64倍直接导致推理速度变慢、成本飙升。信息稀释与噪声干扰并非文档中的所有内容都与当前问题相关。大量无关的段落、细节描述、重复内容会稀释关键信息在模型注意力中的权重相当于在重要的信号中混入了大量噪声导致模型难以聚焦。有限的上下文窗口即便是当前支持较长上下文的模型如128K其有效处理能力也并非与窗口长度线性增长。模型在超长上下文中检索和关联信息的能力会下降容易出现“中间丢失”现象即对位于上下文中部的信息记忆和理解变弱。提示词Prompt结构混乱直接将长文本拼接进Prompt会破坏Prompt的清晰结构使得系统指令、用户问题、参考文档之间的界限模糊影响模型对任务意图的理解。SAGE的设计目标就是要在将文档送入LLM主模型之前先进行一次智能的、有损的“压缩”过滤掉噪声保留精华。2.2 核心思想注意力机制驱动的动态摘要SAGE的核心思想非常直观利用一个轻量级的“压缩器”模型模拟目标LLM在理想情况下会对长文档施加的注意力分布然后根据这个分布来抽取最重要的文本片段。这个过程可以分解为几个关键步骤编码与表示首先将长文档分割成一系列连续的、有重叠的片段chunks。然后使用一个预训练的语言模型编码器例如BERT、RoBERTa为每个片段生成一个高维的向量表示embedding。这个编码器是固定权重的不参与训练只负责将文本转化为数学表示。学习注意力权重这是SAGE的灵魂。框架会训练一个轻量级的“注意力网络”。这个网络的输入是a) 当前用户的问题Query的嵌入表示b) 所有文档片段的嵌入表示。它的任务是预测如果我们将完整的文档和问题一起交给那个庞大的、我们最终要用的目标LLM例如GPT-4、Claude目标LLM的注意力机制会对每个文档片段分配多少“关注度”。基于权重的片段选择根据学习到的注意力权重对所有文档片段进行排序。选择权重最高的前K个片段K是一个可配置的超参数决定了压缩后的上下文长度。重构压缩上下文将被选中的K个片段按照它们在原文档中的顺序拼接起来形成一个新的、缩短了的“压缩文档”。这个压缩文档连同原始的用户问题一起构成最终的Prompt发送给目标LLM进行答案生成。注意这里的关键在于SAGE不是在做传统的“文本摘要”。传统摘要的目标是生成一段连贯、全面的概括性文字。而SAGE的目标是为特定问题筛选出最相关的文本证据。它输出的可能是不连贯的、跳跃的片段集合但只要这些片段包含了回答问题所需的关键信息任务就成功了。这是一种“检索增强”的思想但检索过程是通过模拟注意力来完成的而非简单的关键词匹配。2.3 架构拆解三个核心组件SAGE框架通常包含三个核心组件上下文编码器Context Encoder负责将长文档分割并编码为片段向量。它通常是一个冻结的不训练预训练模型保证文本表示的通用性。注意力预测网络Attention Predictor这是需要训练的核心模块。它是一个轻量级神经网络例如几层MLP或Transformer层输入是问题向量和所有片段向量输出是一个与片段数量相同的权重向量每个权重值在0到1之间代表该片段的重要性。目标LLMTarget LLM即最终用于生成答案的大型语言模型如GPT-4、Llama 3等。SAGE框架本身并不修改这个模型只是优化它的输入。训练SAGE框架需要构建一个数据集其中每个样本包含一个长文档、一个问题、该问题在完整文档下的标准答案或由强大LLM生成的答案。训练时用完整文档输入目标LLM得到答案作为“教师信号”同时记录下目标LLM在处理过程中对各片段的注意力权重或通过某种代理方式估算。然后用这个注意力权重作为标签来训练“注意力预测网络”让它学会根据问题和片段预测出这些权重。3. 关键技术细节与实现要点3.1 文档分块策略重叠与边界处理文档分块是第一步其质量直接影响后续压缩效果。简单的按固定字符数或句子数切割会割裂完整的语义单元。策略推荐使用基于语义的分块工具如LangChain的RecursiveCharacterTextSplitter它可以尝试在段落、句子等自然边界处进行分割。更高级的做法是使用专门的分句和分段模型。重叠Overlap在分块时设置一定的重叠长度例如100-200个字符至关重要。这可以防止一个关键信息恰好被切割在两个块的交界处导致两个块的向量表示都无法完整捕获该信息从而在压缩阶段被遗漏。重叠确保了上下文的连续性。块大小Chunk Size需要权衡。块太小会产生大量片段增加编码和注意力预测的计算量块太大则每个片段包含的信息可能过于混杂不利于精准定位。通常块大小与编码器模型的最大输入长度对齐如BERT的512子词是一个合理的起点。3.2 注意力权重的获取与监督信号构建训练“注意力预测网络”最大的挑战在于如何获得目标LLM对每个文档片段的“真实”注意力权重作为训练标签直接获取商业闭源LLM如GPT-4的内部注意力矩阵几乎不可能。SAGE论文中提出了几种巧妙的代理方法基于输出的重要性评分将完整文档输入目标LLM让其生成答案。然后对于每一个文档片段将其从上下文中移除再次输入目标LLM生成答案。通过比较两次答案的质量例如使用ROUGE分数或答案正确性判断该片段的重要性下降程度就可以作为其注意力权重的代理标签。下降越多说明该片段越重要。基于梯度的显著性方法对于开源模型可以使用集成梯度Integrated Gradients或类似方法计算输入片段中每个token对最终输出答案的贡献度然后将整个片段的贡献度聚合作为该片段的权重。利用模型内部注意力对于完全开源且可访问内部状态的目标LLM可以直接提取其某一层通常是最后几层的注意力矩阵并对与文档片段相关的注意力头进行聚合平均得到片段级权重。在实际操作中方法1基于输出的重要性评分最为通用不依赖于模型内部结构但计算成本最高因为需要为每个训练样本的每个片段都进行一次推理。通常需要采样部分片段进行训练。3.3 注意力预测网络的设计这是一个相对轻量的模块设计目标是高效和有效。输入表示将用户问题编码为一个固定长度的向量q。将每个文档片段编码为向量序列{c1, c2, ..., cn}。通常会对片段向量序列进行池化如平均池化得到每个片段的单一向量表示ci。交互计算计算问题与每个片段的相关性。一种简单有效的方法是计算点积或余弦相似度score_i sim(q, ci)。但SAGE可以做得更复杂例如使用一个交叉注意力模块让问题向量与所有片段向量进行交互生成上下文感知的片段表示。权重归一化将计算出的原始分数scores通过softmax函数归一化为概率分布即weight_i exp(score_i) / sum(exp(score_j))。这确保了所有权重之和为1并且具有可比性。网络结构可以是一个简单的双线性层score_i q^T * W * ci其中W是可学习参数矩阵。也可以是一个小型的Transformer层同时处理所有片段和问题。3.4 压缩率与片段数量K的选择压缩率是SAGE的核心超参数定义为K / N其中N是原始文档总片段数K是选择的片段数。选择策略K的选择没有黄金标准需要根据目标LLM的上下文窗口限制和任务难度进行权衡。固定K例如无论文档多长只选择最重要的10个片段。这能严格保证输入长度但可能对超长文档信息损失过大。自适应K设定一个阈值例如只选择注意力权重累计和达到总权重90%的片段。这能动态调整压缩率保证核心信息不丢失但输出长度不确定。经验值在多项实验中对于万token级别的文档将其压缩到原始长度的20%-30%即保留最相关的20%-30%的文本往往能在保持95%以上问答准确率的同时大幅降低计算成本。这是一个值得参考的起点。4. 实操构建与核心代码解析下面我们将以一个简化的流程演示如何为一个开源LLM这里以Llama 3为例构建一个SAGE风格的上下文压缩管道。我们将使用基于输出的重要性评分方法来生成训练数据。4.1 环境准备与依赖安装首先确保你的Python环境建议3.9以上并安装必要库。我们将使用transformers,torch,langchain用于文本处理openai或litellm作为访问强大LLM生成训练数据的接口假设我们使用GPT-4作为“教师”来生成答案和评估重要性。pip install transformers torch langchain openai tiktoken numpy4.2 步骤一文档预处理与分块我们使用LangChain的文本分割器。from langchain.text_splitter import RecursiveCharacterTextSplitter def chunk_document(full_text, chunk_size500, chunk_overlap100): 将长文档分割成有重叠的块。 Args: full_text: 完整的文档字符串。 chunk_size: 每个块的大致字符数。 chunk_overlap: 块之间的重叠字符数。 Returns: chunks: 文本块列表。 chunk_indices: 每个块在原文中的起止位置列表。 text_splitter RecursiveCharacterTextSplitter( chunk_sizechunk_size, chunk_overlapchunk_overlap, length_functionlen, separators[\n\n, \n, 。, , , , , , ] ) chunks text_splitter.split_text(full_text) # 简化版这里不精确计算indices实际应用需要记录位置信息用于后续重构。 return chunks # 示例 with open(long_document.txt, r, encodingutf-8) as f: doc_text f.read() document_chunks chunk_document(doc_text) print(f文档被分割成 {len(document_chunks)} 个块。)4.3 步骤二构建训练数据模拟注意力标签这是最耗资源的步骤。我们需要为每个文档问题对评估每个块的重要性。import openai import numpy as np client openai.OpenAI(api_keyyour-api-key) # 使用GPT-4作为教师模型 def get_answer_from_llm(context, question, modelgpt-4-turbo): 调用LLM获取答案。 prompt f请基于以下上下文回答问题。如果上下文不包含答案请说“根据上下文无法回答”。 上下文{context} 问题{question} 答案 response client.chat.completions.create( modelmodel, messages[{role: user, content: prompt}], temperature0, max_tokens200 ) return response.choices[0].message.content.strip() def evaluate_answer_quality(answer, reference_answer): 简单评估答案质量。这里使用ROUGE或BERTScore更佳为简化使用字符串包含判断。 # 这是一个极其简化的评估实际应用中应使用更可靠的指标。 # 假设我们有标准答案这里用参考答案模拟。 # 真实情况可能是用GPT-4自己判断两次答案的优劣。 return 1.0 if reference_answer in answer else 0.0 def generate_training_sample(full_doc_chunks, question, reference_answer): 为一个文档问题对生成训练样本。 返回每个chunk的注意力权重标签。 n_chunks len(full_doc_chunks) weights np.zeros(n_chunks) # 1. 获取完整上下文下的答案作为基准 full_context \n\n.join(full_doc_chunks) baseline_answer get_answer_from_llm(full_context, question) baseline_score evaluate_answer_quality(baseline_answer, reference_answer) # 2. 遍历每个chunk将其移除后评估答案质量下降程度 for i in range(n_chunks): ablated_chunks full_doc_chunks[:i] full_doc_chunks[i1:] ablated_context \n\n.join(ablated_chunks) ablated_answer get_answer_from_llm(ablated_context, question) ablated_score evaluate_answer_quality(ablated_answer, reference_answer) # 质量下降幅度作为该chunk重要性的近似 importance max(0, baseline_score - ablated_score) # 确保非负 weights[i] importance # 3. 归一化权重使其和为1模拟注意力分布 if weights.sum() 0: weights weights / weights.sum() else: weights np.ones(n_chunks) / n_chunks # 如果所有块都不重要则均匀分布 return weights # 注意此函数调用API次数为 O(N1)成本高仅用于演示原理。 # 实际研究中会采用采样、缓存等策略优化。4.4 步骤三实现注意力预测网络我们实现一个简单的双线性注意力网络。import torch import torch.nn as nn import torch.nn.functional as F class SageAttentionPredictor(nn.Module): def __init__(self, embedding_dim768): super().__init__() # 将问题和文档片段映射到同一空间并计算相关性 self.query_proj nn.Linear(embedding_dim, embedding_dim) self.key_proj nn.Linear(embedding_dim, embedding_dim) # 可选的一个小的融合层 self.fusion nn.Sequential( nn.Linear(embedding_dim * 2, embedding_dim), nn.ReLU(), nn.Linear(embedding_dim, 1) ) def forward(self, question_embedding, chunk_embeddings): question_embedding: [batch_size, embed_dim] chunk_embeddings: [batch_size, num_chunks, embed_dim] 返回: [batch_size, num_chunks] 注意力权重 batch_size, num_chunks, _ chunk_embeddings.shape # 投影 q self.query_proj(question_embedding).unsqueeze(1) # [batch, 1, dim] k self.key_proj(chunk_embeddings) # [batch, num_chunks, dim] # 方法1简单点积注意力 # scores torch.bmm(q, k.transpose(1, 2)).squeeze(1) # [batch, num_chunks] # 方法2使用融合网络更灵活 q_expanded q.expand(-1, num_chunks, -1) # [batch, num_chunks, dim] combined torch.cat([q_expanded, k], dim-1) # [batch, num_chunks, dim*2] scores self.fusion(combined).squeeze(-1) # [batch, num_chunks] # 归一化为概率分布 weights F.softmax(scores, dim-1) return weights4.5 步骤四训练与推理管道训练循环和推理管道的简化示例。# 假设我们有数据集list of (question_embedding, chunk_embeddings, target_weights) # question_embedding: 通过BERT等编码器编码问题得到 # chunk_embeddings: 通过同一个编码器编码所有文档块得到 # target_weights: 上一步generate_training_sample生成的权重 def train_sage(model, train_loader, epochs10): optimizer torch.optim.Adam(model.parameters(), lr1e-4) criterion nn.KLDivLoss(reductionbatchmean) # 使用KL散度损失因为目标是分布 model.train() for epoch in range(epochs): total_loss 0 for q_emb, c_emb, target_w in train_loader: optimizer.zero_grad() pred_weights model(q_emb, c_emb) # 预测权重和真实权重都是概率分布 loss criterion(pred_weights.log(), target_w) # KLDivLoss需要log输入 loss.backward() optimizer.step() total_loss loss.item() print(fEpoch {epoch1}, Loss: {total_loss/len(train_loader):.4f}) def compress_context_with_sage(model, question, document_chunks, encoder_model, top_k5): 使用训练好的SAGE模型压缩上下文。 # 1. 编码 with torch.no_grad(): q_embedding encoder_model.encode(question, convert_to_tensorTrue).unsqueeze(0) # [1, dim] chunk_embeddings [] for chunk in document_chunks: emb encoder_model.encode(chunk, convert_to_tensorTrue) chunk_embeddings.append(emb) chunk_embeddings torch.stack(chunk_embeddings).unsqueeze(0) # [1, num_chunks, dim] # 2. 预测注意力权重 attention_weights model(q_embedding, chunk_embeddings) # [1, num_chunks] weights attention_weights.squeeze(0).cpu().numpy() # 3. 选择top-k个块 top_indices np.argsort(weights)[-top_k:][::-1] # 从高到低排序 selected_chunks [document_chunks[i] for i in top_indices] # 4. 按原顺序排序并拼接 selected_indices_sorted sorted(top_indices) final_context \n\n.join([document_chunks[i] for i in selected_indices_sorted]) return final_context, weights, selected_indices_sorted # 压缩后的final_context即可送入目标LLM进行问答。5. 常见问题、挑战与优化策略在实际部署SAGE或类似框架时你会遇到一系列工程和算法上的挑战。5.1 训练数据获取成本高昂问题使用基于输出的重要性评分方法每个训练样本需要对每个文档块进行一次LLM推理成本随文档长度线性增长几乎不可行。解决方案采样训练不对所有块进行评估而是随机采样一部分块例如20%进行重要性评估其余块权重设为0或一个很小的基线值。这能大幅降低成本。合成数据生成利用已有的短上下文QA数据集如SQuAD通过“上下文膨胀”技术人工添加无关的干扰段落构建“长文档”和对应的“压缩后核心段落”。这样可以低成本获得大量训练数据。知识蒸馏先用成本高的方法如GPT-4评估在一个小型高质量数据集上训练一个“教师”SAGE模型然后用这个教师模型为更大的、未标注的长文档数据集生成注意力权重标签再用这些标签训练一个更轻量的“学生”模型。5.2 注意力预测的泛化能力问题在一个领域如医学文献上训练的SAGE模型在另一个领域如法律文书上表现可能下降。解决方案领域适配在目标领域的小规模数据上进行微调Fine-tuning。即使只有几百个样本也能显著提升性能。使用领域无关的编码器采用在多种语料上预训练的通用文本编码器如all-MiniLM-L6-v2增强模型的泛化基础。多任务学习在训练时除了预测注意力权重还可以让模型同时学习一些辅助任务如句子相似度判断、下一句预测等以学习更通用的文本表示。5.3 压缩导致的上下文断裂问题选出的Top-K个片段可能在原文中不相邻直接拼接会破坏局部的连贯性可能影响目标LLM对某些需要跨片段推理的理解。解决方案上下文窗口扩展在选中一个片段时同时将其前后相邻的1-2个片段也包含进来即使它们的权重不是最高。这牺牲了一点压缩率但保留了局部上下文。重排序与平滑在按权重选出片段后不是简单按原序拼接而是根据片段间的语义连贯性进行微调或插入简短的连接说明如“接上文...”但这会引入新的文本可能干扰模型。两阶段压缩第一阶段用SAGE选出候选片段池数量稍多于K第二阶段用一个轻量级模型对这些候选片段进行连贯性重排和微调再选择最连贯的子集。5.4 实时推理延迟问题SAGE压缩过程本身需要编码文档块和运行预测网络这会为问答流程增加额外的延迟。优化策略预编码与缓存对于静态或更新不频繁的长文档如知识库文章可以预先计算所有文档块的向量嵌入并缓存。当新问题到来时只需编码问题并计算注意力大大减少延迟。模型轻量化将注意力预测网络设计得尽可能小使用蒸馏后的模型。编码器也可以选择更快的模型如Sentence-Transformers中的轻量级模型。异步处理在用户提问前对热门或预期的文档进行预压缩处理。5.5 评估指标的选择问题如何量化SAGE框架的有效性不能只看压缩率。核心评估维度任务性能压缩后的上下文输入目标LLM其问答的准确率EM、F1分数、ROUGE分数等与使用完整上下文相比的保留百分比。这是最重要的指标。效率提升推理速度的加快比例、Token消耗的减少比例、成本的下降比例。压缩质量选中的片段是否真的包含答案依据可以通过人工评估或与标准证据片段的重合度如召回率来衡量。一个实用的评估流程是在测试集上分别用“完整上下文”、“随机选择K个片段”、“SAGE选择K个片段”三种方式输入目标LLM比较三者的任务性能。SAGE应该在性能接近完整上下文的同时显著优于随机选择。6. 进阶应用与未来展望SAGE的思想可以扩展到更广泛的场景不局限于简单的抽取式压缩。迭代式压缩与交互在多轮对话中第一轮用SAGE压缩原始文档得到初始上下文并生成回答。如果用户追问细节可以将上一轮选中的片段及其相邻区域作为新的“长文档”结合新问题再次运行SAGE进行二次压缩实现动态的、聚焦式的上下文管理。与检索增强生成RAG的融合在经典的RAG流程中先用向量数据库检索出Top-K个相关片段。可以在此基础上引入SAGE机制将这K个检索结果视为“候选片段”再用一个轻量的注意力预测网络根据具体问题对这K个片段进行重要性重排和进一步筛选只将最重要的几个送入LLM进一步提升效率和质量。多模态扩展对于包含图文、表格的长文档SAGE的思想同样适用。需要设计能够处理多模态信息的编码器和注意力预测网络实现对图像区域、表格单元格等重要度的评估和选择。无监督/自监督学习探索不依赖昂贵人工或LLM标注的注意力权重获取方法。例如利用文档自身的结构信息标题、加粗、引用、或基于信息论的方法如某个片段能多大程度降低问题答案的不确定性来生成训练信号。在实际操作中完全复现论文中的SAGE可能需要庞大的计算资源进行训练。但对于大多数应用者来说理解其原理后可以采用更工程化的简化方案例如用双编码器Dual Encoder计算问题与片段的相似度作为权重或者直接用现成的交叉编码器Cross-Encoder对每个问题片段对进行打分排序。这些方法虽然不如端到端训练的SAGE精准但实现简单在特定场景下也能获得显著收益。最终SAGE框架给我们最大的启示是面对长上下文挑战与其一味追求扩大模型的窗口不如思考如何更智能地管理输入的信息。让一个小而专的模型去学习“什么是重要的”然后把最重要的部分交给大模型去深度处理这种分工协作的思路在AI系统设计中会越来越普遍。
SAGE框架：基于注意力机制的长文档问答上下文压缩技术解析
1. 项目概述当长文档问答遇上“信息过载”处理一份动辄几十页、上百页的PDF报告或者一本电子书然后向AI模型提出一个具体问题这可能是很多研究者和开发者正在尝试的事情。理想很丰满我们把整个文档“喂”给模型它就能像一位通读了全文的专家给出精准的答案。但现实很骨感当你真的把一篇数万token的文档塞进提示词Prompt时往往会发现模型的表现不尽如人意——它可能答非所问遗漏关键细节或者干脆因为上下文长度限制而拒绝处理。核心矛盾就在于我们提供给模型的“上下文窗口”是有限的而长文档的信息是冗余且分散的。这就是“SAGE基于注意力机制的上下文压缩框架”要解决的核心问题。它不是一个全新的模型而是一个精巧的“预处理”框架。你可以把它想象成一位高效的“研究助理”在将长文档交给大语言模型LLM这位“主分析师”之前这位助理会先快速浏览全文然后不是原封不动地递交所有材料而是整理出一份高度凝练、只包含与当前问题最相关信息的“摘要简报”。SAGE的创新之处在于它借鉴了Transformer模型核心的“注意力机制”思想来模拟这个“快速浏览并抓取重点”的过程从而实现自适应的、动态的上下文压缩。对于需要处理长文本问答Long-Document QA的开发者、知识库构建者或是任何受限于模型上下文长度但又要保证回答质量的应用场景理解SAGE的原理和实现思路远比简单地调用一个API更有价值。它能帮你从根本上优化信息流用更低的计算和成本开销换取更精准、更可靠的问答效果。2. SAGE框架的核心设计思想与工作原理2.1 问题根源为什么直接输入长文档会失效要理解SAGE的价值首先得明白为什么“暴力”输入长文档效果不好。这背后有几个关键原因计算复杂度与成本Transformer模型的自注意力机制的计算复杂度与序列长度的平方成正比。这意味着将上下文长度从1K扩展到8K计算负担可能增加64倍直接导致推理速度变慢、成本飙升。信息稀释与噪声干扰并非文档中的所有内容都与当前问题相关。大量无关的段落、细节描述、重复内容会稀释关键信息在模型注意力中的权重相当于在重要的信号中混入了大量噪声导致模型难以聚焦。有限的上下文窗口即便是当前支持较长上下文的模型如128K其有效处理能力也并非与窗口长度线性增长。模型在超长上下文中检索和关联信息的能力会下降容易出现“中间丢失”现象即对位于上下文中部的信息记忆和理解变弱。提示词Prompt结构混乱直接将长文本拼接进Prompt会破坏Prompt的清晰结构使得系统指令、用户问题、参考文档之间的界限模糊影响模型对任务意图的理解。SAGE的设计目标就是要在将文档送入LLM主模型之前先进行一次智能的、有损的“压缩”过滤掉噪声保留精华。2.2 核心思想注意力机制驱动的动态摘要SAGE的核心思想非常直观利用一个轻量级的“压缩器”模型模拟目标LLM在理想情况下会对长文档施加的注意力分布然后根据这个分布来抽取最重要的文本片段。这个过程可以分解为几个关键步骤编码与表示首先将长文档分割成一系列连续的、有重叠的片段chunks。然后使用一个预训练的语言模型编码器例如BERT、RoBERTa为每个片段生成一个高维的向量表示embedding。这个编码器是固定权重的不参与训练只负责将文本转化为数学表示。学习注意力权重这是SAGE的灵魂。框架会训练一个轻量级的“注意力网络”。这个网络的输入是a) 当前用户的问题Query的嵌入表示b) 所有文档片段的嵌入表示。它的任务是预测如果我们将完整的文档和问题一起交给那个庞大的、我们最终要用的目标LLM例如GPT-4、Claude目标LLM的注意力机制会对每个文档片段分配多少“关注度”。基于权重的片段选择根据学习到的注意力权重对所有文档片段进行排序。选择权重最高的前K个片段K是一个可配置的超参数决定了压缩后的上下文长度。重构压缩上下文将被选中的K个片段按照它们在原文档中的顺序拼接起来形成一个新的、缩短了的“压缩文档”。这个压缩文档连同原始的用户问题一起构成最终的Prompt发送给目标LLM进行答案生成。注意这里的关键在于SAGE不是在做传统的“文本摘要”。传统摘要的目标是生成一段连贯、全面的概括性文字。而SAGE的目标是为特定问题筛选出最相关的文本证据。它输出的可能是不连贯的、跳跃的片段集合但只要这些片段包含了回答问题所需的关键信息任务就成功了。这是一种“检索增强”的思想但检索过程是通过模拟注意力来完成的而非简单的关键词匹配。2.3 架构拆解三个核心组件SAGE框架通常包含三个核心组件上下文编码器Context Encoder负责将长文档分割并编码为片段向量。它通常是一个冻结的不训练预训练模型保证文本表示的通用性。注意力预测网络Attention Predictor这是需要训练的核心模块。它是一个轻量级神经网络例如几层MLP或Transformer层输入是问题向量和所有片段向量输出是一个与片段数量相同的权重向量每个权重值在0到1之间代表该片段的重要性。目标LLMTarget LLM即最终用于生成答案的大型语言模型如GPT-4、Llama 3等。SAGE框架本身并不修改这个模型只是优化它的输入。训练SAGE框架需要构建一个数据集其中每个样本包含一个长文档、一个问题、该问题在完整文档下的标准答案或由强大LLM生成的答案。训练时用完整文档输入目标LLM得到答案作为“教师信号”同时记录下目标LLM在处理过程中对各片段的注意力权重或通过某种代理方式估算。然后用这个注意力权重作为标签来训练“注意力预测网络”让它学会根据问题和片段预测出这些权重。3. 关键技术细节与实现要点3.1 文档分块策略重叠与边界处理文档分块是第一步其质量直接影响后续压缩效果。简单的按固定字符数或句子数切割会割裂完整的语义单元。策略推荐使用基于语义的分块工具如LangChain的RecursiveCharacterTextSplitter它可以尝试在段落、句子等自然边界处进行分割。更高级的做法是使用专门的分句和分段模型。重叠Overlap在分块时设置一定的重叠长度例如100-200个字符至关重要。这可以防止一个关键信息恰好被切割在两个块的交界处导致两个块的向量表示都无法完整捕获该信息从而在压缩阶段被遗漏。重叠确保了上下文的连续性。块大小Chunk Size需要权衡。块太小会产生大量片段增加编码和注意力预测的计算量块太大则每个片段包含的信息可能过于混杂不利于精准定位。通常块大小与编码器模型的最大输入长度对齐如BERT的512子词是一个合理的起点。3.2 注意力权重的获取与监督信号构建训练“注意力预测网络”最大的挑战在于如何获得目标LLM对每个文档片段的“真实”注意力权重作为训练标签直接获取商业闭源LLM如GPT-4的内部注意力矩阵几乎不可能。SAGE论文中提出了几种巧妙的代理方法基于输出的重要性评分将完整文档输入目标LLM让其生成答案。然后对于每一个文档片段将其从上下文中移除再次输入目标LLM生成答案。通过比较两次答案的质量例如使用ROUGE分数或答案正确性判断该片段的重要性下降程度就可以作为其注意力权重的代理标签。下降越多说明该片段越重要。基于梯度的显著性方法对于开源模型可以使用集成梯度Integrated Gradients或类似方法计算输入片段中每个token对最终输出答案的贡献度然后将整个片段的贡献度聚合作为该片段的权重。利用模型内部注意力对于完全开源且可访问内部状态的目标LLM可以直接提取其某一层通常是最后几层的注意力矩阵并对与文档片段相关的注意力头进行聚合平均得到片段级权重。在实际操作中方法1基于输出的重要性评分最为通用不依赖于模型内部结构但计算成本最高因为需要为每个训练样本的每个片段都进行一次推理。通常需要采样部分片段进行训练。3.3 注意力预测网络的设计这是一个相对轻量的模块设计目标是高效和有效。输入表示将用户问题编码为一个固定长度的向量q。将每个文档片段编码为向量序列{c1, c2, ..., cn}。通常会对片段向量序列进行池化如平均池化得到每个片段的单一向量表示ci。交互计算计算问题与每个片段的相关性。一种简单有效的方法是计算点积或余弦相似度score_i sim(q, ci)。但SAGE可以做得更复杂例如使用一个交叉注意力模块让问题向量与所有片段向量进行交互生成上下文感知的片段表示。权重归一化将计算出的原始分数scores通过softmax函数归一化为概率分布即weight_i exp(score_i) / sum(exp(score_j))。这确保了所有权重之和为1并且具有可比性。网络结构可以是一个简单的双线性层score_i q^T * W * ci其中W是可学习参数矩阵。也可以是一个小型的Transformer层同时处理所有片段和问题。3.4 压缩率与片段数量K的选择压缩率是SAGE的核心超参数定义为K / N其中N是原始文档总片段数K是选择的片段数。选择策略K的选择没有黄金标准需要根据目标LLM的上下文窗口限制和任务难度进行权衡。固定K例如无论文档多长只选择最重要的10个片段。这能严格保证输入长度但可能对超长文档信息损失过大。自适应K设定一个阈值例如只选择注意力权重累计和达到总权重90%的片段。这能动态调整压缩率保证核心信息不丢失但输出长度不确定。经验值在多项实验中对于万token级别的文档将其压缩到原始长度的20%-30%即保留最相关的20%-30%的文本往往能在保持95%以上问答准确率的同时大幅降低计算成本。这是一个值得参考的起点。4. 实操构建与核心代码解析下面我们将以一个简化的流程演示如何为一个开源LLM这里以Llama 3为例构建一个SAGE风格的上下文压缩管道。我们将使用基于输出的重要性评分方法来生成训练数据。4.1 环境准备与依赖安装首先确保你的Python环境建议3.9以上并安装必要库。我们将使用transformers,torch,langchain用于文本处理openai或litellm作为访问强大LLM生成训练数据的接口假设我们使用GPT-4作为“教师”来生成答案和评估重要性。pip install transformers torch langchain openai tiktoken numpy4.2 步骤一文档预处理与分块我们使用LangChain的文本分割器。from langchain.text_splitter import RecursiveCharacterTextSplitter def chunk_document(full_text, chunk_size500, chunk_overlap100): 将长文档分割成有重叠的块。 Args: full_text: 完整的文档字符串。 chunk_size: 每个块的大致字符数。 chunk_overlap: 块之间的重叠字符数。 Returns: chunks: 文本块列表。 chunk_indices: 每个块在原文中的起止位置列表。 text_splitter RecursiveCharacterTextSplitter( chunk_sizechunk_size, chunk_overlapchunk_overlap, length_functionlen, separators[\n\n, \n, 。, , , , , , ] ) chunks text_splitter.split_text(full_text) # 简化版这里不精确计算indices实际应用需要记录位置信息用于后续重构。 return chunks # 示例 with open(long_document.txt, r, encodingutf-8) as f: doc_text f.read() document_chunks chunk_document(doc_text) print(f文档被分割成 {len(document_chunks)} 个块。)4.3 步骤二构建训练数据模拟注意力标签这是最耗资源的步骤。我们需要为每个文档问题对评估每个块的重要性。import openai import numpy as np client openai.OpenAI(api_keyyour-api-key) # 使用GPT-4作为教师模型 def get_answer_from_llm(context, question, modelgpt-4-turbo): 调用LLM获取答案。 prompt f请基于以下上下文回答问题。如果上下文不包含答案请说“根据上下文无法回答”。 上下文{context} 问题{question} 答案 response client.chat.completions.create( modelmodel, messages[{role: user, content: prompt}], temperature0, max_tokens200 ) return response.choices[0].message.content.strip() def evaluate_answer_quality(answer, reference_answer): 简单评估答案质量。这里使用ROUGE或BERTScore更佳为简化使用字符串包含判断。 # 这是一个极其简化的评估实际应用中应使用更可靠的指标。 # 假设我们有标准答案这里用参考答案模拟。 # 真实情况可能是用GPT-4自己判断两次答案的优劣。 return 1.0 if reference_answer in answer else 0.0 def generate_training_sample(full_doc_chunks, question, reference_answer): 为一个文档问题对生成训练样本。 返回每个chunk的注意力权重标签。 n_chunks len(full_doc_chunks) weights np.zeros(n_chunks) # 1. 获取完整上下文下的答案作为基准 full_context \n\n.join(full_doc_chunks) baseline_answer get_answer_from_llm(full_context, question) baseline_score evaluate_answer_quality(baseline_answer, reference_answer) # 2. 遍历每个chunk将其移除后评估答案质量下降程度 for i in range(n_chunks): ablated_chunks full_doc_chunks[:i] full_doc_chunks[i1:] ablated_context \n\n.join(ablated_chunks) ablated_answer get_answer_from_llm(ablated_context, question) ablated_score evaluate_answer_quality(ablated_answer, reference_answer) # 质量下降幅度作为该chunk重要性的近似 importance max(0, baseline_score - ablated_score) # 确保非负 weights[i] importance # 3. 归一化权重使其和为1模拟注意力分布 if weights.sum() 0: weights weights / weights.sum() else: weights np.ones(n_chunks) / n_chunks # 如果所有块都不重要则均匀分布 return weights # 注意此函数调用API次数为 O(N1)成本高仅用于演示原理。 # 实际研究中会采用采样、缓存等策略优化。4.4 步骤三实现注意力预测网络我们实现一个简单的双线性注意力网络。import torch import torch.nn as nn import torch.nn.functional as F class SageAttentionPredictor(nn.Module): def __init__(self, embedding_dim768): super().__init__() # 将问题和文档片段映射到同一空间并计算相关性 self.query_proj nn.Linear(embedding_dim, embedding_dim) self.key_proj nn.Linear(embedding_dim, embedding_dim) # 可选的一个小的融合层 self.fusion nn.Sequential( nn.Linear(embedding_dim * 2, embedding_dim), nn.ReLU(), nn.Linear(embedding_dim, 1) ) def forward(self, question_embedding, chunk_embeddings): question_embedding: [batch_size, embed_dim] chunk_embeddings: [batch_size, num_chunks, embed_dim] 返回: [batch_size, num_chunks] 注意力权重 batch_size, num_chunks, _ chunk_embeddings.shape # 投影 q self.query_proj(question_embedding).unsqueeze(1) # [batch, 1, dim] k self.key_proj(chunk_embeddings) # [batch, num_chunks, dim] # 方法1简单点积注意力 # scores torch.bmm(q, k.transpose(1, 2)).squeeze(1) # [batch, num_chunks] # 方法2使用融合网络更灵活 q_expanded q.expand(-1, num_chunks, -1) # [batch, num_chunks, dim] combined torch.cat([q_expanded, k], dim-1) # [batch, num_chunks, dim*2] scores self.fusion(combined).squeeze(-1) # [batch, num_chunks] # 归一化为概率分布 weights F.softmax(scores, dim-1) return weights4.5 步骤四训练与推理管道训练循环和推理管道的简化示例。# 假设我们有数据集list of (question_embedding, chunk_embeddings, target_weights) # question_embedding: 通过BERT等编码器编码问题得到 # chunk_embeddings: 通过同一个编码器编码所有文档块得到 # target_weights: 上一步generate_training_sample生成的权重 def train_sage(model, train_loader, epochs10): optimizer torch.optim.Adam(model.parameters(), lr1e-4) criterion nn.KLDivLoss(reductionbatchmean) # 使用KL散度损失因为目标是分布 model.train() for epoch in range(epochs): total_loss 0 for q_emb, c_emb, target_w in train_loader: optimizer.zero_grad() pred_weights model(q_emb, c_emb) # 预测权重和真实权重都是概率分布 loss criterion(pred_weights.log(), target_w) # KLDivLoss需要log输入 loss.backward() optimizer.step() total_loss loss.item() print(fEpoch {epoch1}, Loss: {total_loss/len(train_loader):.4f}) def compress_context_with_sage(model, question, document_chunks, encoder_model, top_k5): 使用训练好的SAGE模型压缩上下文。 # 1. 编码 with torch.no_grad(): q_embedding encoder_model.encode(question, convert_to_tensorTrue).unsqueeze(0) # [1, dim] chunk_embeddings [] for chunk in document_chunks: emb encoder_model.encode(chunk, convert_to_tensorTrue) chunk_embeddings.append(emb) chunk_embeddings torch.stack(chunk_embeddings).unsqueeze(0) # [1, num_chunks, dim] # 2. 预测注意力权重 attention_weights model(q_embedding, chunk_embeddings) # [1, num_chunks] weights attention_weights.squeeze(0).cpu().numpy() # 3. 选择top-k个块 top_indices np.argsort(weights)[-top_k:][::-1] # 从高到低排序 selected_chunks [document_chunks[i] for i in top_indices] # 4. 按原顺序排序并拼接 selected_indices_sorted sorted(top_indices) final_context \n\n.join([document_chunks[i] for i in selected_indices_sorted]) return final_context, weights, selected_indices_sorted # 压缩后的final_context即可送入目标LLM进行问答。5. 常见问题、挑战与优化策略在实际部署SAGE或类似框架时你会遇到一系列工程和算法上的挑战。5.1 训练数据获取成本高昂问题使用基于输出的重要性评分方法每个训练样本需要对每个文档块进行一次LLM推理成本随文档长度线性增长几乎不可行。解决方案采样训练不对所有块进行评估而是随机采样一部分块例如20%进行重要性评估其余块权重设为0或一个很小的基线值。这能大幅降低成本。合成数据生成利用已有的短上下文QA数据集如SQuAD通过“上下文膨胀”技术人工添加无关的干扰段落构建“长文档”和对应的“压缩后核心段落”。这样可以低成本获得大量训练数据。知识蒸馏先用成本高的方法如GPT-4评估在一个小型高质量数据集上训练一个“教师”SAGE模型然后用这个教师模型为更大的、未标注的长文档数据集生成注意力权重标签再用这些标签训练一个更轻量的“学生”模型。5.2 注意力预测的泛化能力问题在一个领域如医学文献上训练的SAGE模型在另一个领域如法律文书上表现可能下降。解决方案领域适配在目标领域的小规模数据上进行微调Fine-tuning。即使只有几百个样本也能显著提升性能。使用领域无关的编码器采用在多种语料上预训练的通用文本编码器如all-MiniLM-L6-v2增强模型的泛化基础。多任务学习在训练时除了预测注意力权重还可以让模型同时学习一些辅助任务如句子相似度判断、下一句预测等以学习更通用的文本表示。5.3 压缩导致的上下文断裂问题选出的Top-K个片段可能在原文中不相邻直接拼接会破坏局部的连贯性可能影响目标LLM对某些需要跨片段推理的理解。解决方案上下文窗口扩展在选中一个片段时同时将其前后相邻的1-2个片段也包含进来即使它们的权重不是最高。这牺牲了一点压缩率但保留了局部上下文。重排序与平滑在按权重选出片段后不是简单按原序拼接而是根据片段间的语义连贯性进行微调或插入简短的连接说明如“接上文...”但这会引入新的文本可能干扰模型。两阶段压缩第一阶段用SAGE选出候选片段池数量稍多于K第二阶段用一个轻量级模型对这些候选片段进行连贯性重排和微调再选择最连贯的子集。5.4 实时推理延迟问题SAGE压缩过程本身需要编码文档块和运行预测网络这会为问答流程增加额外的延迟。优化策略预编码与缓存对于静态或更新不频繁的长文档如知识库文章可以预先计算所有文档块的向量嵌入并缓存。当新问题到来时只需编码问题并计算注意力大大减少延迟。模型轻量化将注意力预测网络设计得尽可能小使用蒸馏后的模型。编码器也可以选择更快的模型如Sentence-Transformers中的轻量级模型。异步处理在用户提问前对热门或预期的文档进行预压缩处理。5.5 评估指标的选择问题如何量化SAGE框架的有效性不能只看压缩率。核心评估维度任务性能压缩后的上下文输入目标LLM其问答的准确率EM、F1分数、ROUGE分数等与使用完整上下文相比的保留百分比。这是最重要的指标。效率提升推理速度的加快比例、Token消耗的减少比例、成本的下降比例。压缩质量选中的片段是否真的包含答案依据可以通过人工评估或与标准证据片段的重合度如召回率来衡量。一个实用的评估流程是在测试集上分别用“完整上下文”、“随机选择K个片段”、“SAGE选择K个片段”三种方式输入目标LLM比较三者的任务性能。SAGE应该在性能接近完整上下文的同时显著优于随机选择。6. 进阶应用与未来展望SAGE的思想可以扩展到更广泛的场景不局限于简单的抽取式压缩。迭代式压缩与交互在多轮对话中第一轮用SAGE压缩原始文档得到初始上下文并生成回答。如果用户追问细节可以将上一轮选中的片段及其相邻区域作为新的“长文档”结合新问题再次运行SAGE进行二次压缩实现动态的、聚焦式的上下文管理。与检索增强生成RAG的融合在经典的RAG流程中先用向量数据库检索出Top-K个相关片段。可以在此基础上引入SAGE机制将这K个检索结果视为“候选片段”再用一个轻量的注意力预测网络根据具体问题对这K个片段进行重要性重排和进一步筛选只将最重要的几个送入LLM进一步提升效率和质量。多模态扩展对于包含图文、表格的长文档SAGE的思想同样适用。需要设计能够处理多模态信息的编码器和注意力预测网络实现对图像区域、表格单元格等重要度的评估和选择。无监督/自监督学习探索不依赖昂贵人工或LLM标注的注意力权重获取方法。例如利用文档自身的结构信息标题、加粗、引用、或基于信息论的方法如某个片段能多大程度降低问题答案的不确定性来生成训练信号。在实际操作中完全复现论文中的SAGE可能需要庞大的计算资源进行训练。但对于大多数应用者来说理解其原理后可以采用更工程化的简化方案例如用双编码器Dual Encoder计算问题与片段的相似度作为权重或者直接用现成的交叉编码器Cross-Encoder对每个问题片段对进行打分排序。这些方法虽然不如端到端训练的SAGE精准但实现简单在特定场景下也能获得显著收益。最终SAGE框架给我们最大的启示是面对长上下文挑战与其一味追求扩大模型的窗口不如思考如何更智能地管理输入的信息。让一个小而专的模型去学习“什么是重要的”然后把最重要的部分交给大模型去深度处理这种分工协作的思路在AI系统设计中会越来越普遍。
