论文

SproutRAG: 注意力引导的树搜索与渐进式嵌入用于长文档RAG

SproutRAG: 注意力引导的树搜索与渐进式嵌入用于长文档RAG

检索增强生成(RAG)系统需要在检索粒度与上下文连贯性之间取得平衡。现有方法通过LLM引导分块、单级上下文扩展或分层摘要来解决,但这些方法在索引或检索时依赖昂贵的LLM调用、将上下文聚合限制在单一粒度级别,或通过摘要导致信息损失。 我们提出SproutRAG,一个注意力引导的分层RAG框架。它利用学习到的句间注意力将句子级块组织成逐渐增大但语义连贯的单元,构建二元分块树。与依赖外部LLM、固定上下文扩展或有损摘要的方法不同,SproutRAG学习哪些注意力头和层最能捕捉语义文档结构,实现多粒度检索而无需额外LLM调用或压缩摘要。检索时,SproutRAG使用分层束搜索在多个粒度上检索候选,捕获平面检索无法实现的多句相关性。该框架通过联合目标进行端到端训练,同时改进嵌入和树结构。 在涵盖科学、法律和开放域场景的四个基准测试中,SproutRAG相比最强基线平均提高信息效率(IE)6.1%。代码已开源。

论文精读

TL;DR SproutRAG 利用注意力引导的句子级分块构建层次二叉树,实现多粒度检索,端到端训练优化嵌入与树结构,无需外部 LLM 或摘要,信息效率平均提升 6.1%。

问题

问题背景

RAG 系统处理长文档时,需在检索粒度(sentence-level precision)与上下文连贯性(multi-sentence coherence)之间取得平衡。过于精细的切分易丢失语境,过粗的 chunk 则引入噪声。

现有方法局限

当前主流方案各有明显短板:

  • LLM 驱动分块(如 Meta-Chunking、MoC):索引或检索阶段依赖昂贵的外部 LLM 调用,成本与延迟高,难以规模化。
  • 固定窗口扩展(如 Dense X Retrieval):仅支持单一粒度的上下文聚合,无法同时获取多粒度信息,导致跨段落证据检索失败。
  • 层次化摘要(如 RAPTOR):虽构建多级树结构,但以压缩摘要代替原始文本,信息损失不可避免,且无法端到端联合优化嵌入与结构。

这些方法要么牺牲效率,要么损失信息保真度,均未能实现无损多粒度检索无额外 LLM 依赖的兼顾。

为什么难 / 重要

长文档检索的根本难点在于:相关答案常分布在跨越句子和段落的多个位置,模型需同时捕捉细粒度匹配与宽域语境。引入 LLM 虽能缓解,但索引或查询时的额外调用使系统难以应对实时服务与大规模语料库。因此,业界迫切需要一种多粒度感知、无需 LLM 额外介入且保留完整原始文本的检索框架,以提升信息效率和实用性。

行业类比

这类似于法律合同审查中需快速定位条款同时不丢失完整上下文,或研发团队在技术文档中搜索方案时必须兼顾片段精确度和整节逻辑的连贯性。

核心洞察

  • 利用预训练 Transformer 内部的注意力模式发现文档语义结构,无需外部 LLM 辅助。与 Meta-Chunking 等依赖 LLM 指导分块的方法不同,SproutRAG 直接学习哪些注意力头和层最能捕获句子间连贯性,从而构建层次化分块树。这本质上是一种自监督的结构蒸馏,避免了昂贵的 LLM 调用和摘要带来的信息损失,同时保留了原始文本的完整性,为低资源或延迟敏感场景提供了高效索引路径。
  • 多粒度层次化检索与 beam search 的组合,打破了 flat 检索的单一上下文限制。传统方法如 RAPTOR 通过生成摘要构建层次,但信息压缩不可逆;SproutRAG 基于二进制树保留原始句子,在检索时执行层次化 beam search,从粗粒度到细粒度动态探索候选集,有效捕获跨句子长距离依赖。这种结构使模型无须拼凑片段即可获取多句子证据,在信息效率(IE)指标上带来平均 6.1% 的显著提升。

方法

输入与索引构建

SproutRAG 以长文档为输入,首先将文档切分为句子级最小单位 (sentence-level chunks),每个句子用预训练语言模型获得初始 embedding。索引阶段通过 Attention-Guided Indexing 模块,利用模型内部的 inter-sentence attention 构建一个自底向上的二叉合并树 (binary chunking tree)。具体做法是:选择最能反映语义结构的注意力头和层,计算句子之间的关联强度,将语义最邻近的句子对逐层合并为更大的 chunk,并为每个新节点计算渐进式 embedding (progressive embeddings)。整个过程仅依赖模型自身的注意力,无需额外 LLM 调用或生成压缩摘要,从而避免了信息损失。

层次化检索与输出

在检索阶段,SproutRAG 使用 Hierarchical Retrieval 模块,在已构建的树结构上执行层次化 beam search。与传统扁平检索不同,beam search 同时探索不同粒度(句子、段落、文档片段)的节点,保留 top-k 候选路径,从而捕获跨句子的深层语义关联,而不仅仅是单一句子的关键词匹配。最终输出一个多粒度候选集合,包含从细粒度句子到粗粒度段落的多层证据,直接提供给下游生成模型。

端到端联合训练

框架采用端到端联合训练,同时优化两个目标:

  • 检索损失 (retrieval objective):鼓励相关节点在 beam search 中排名更靠前,以提升 top-k 召回质量。
  • 结构损失 (structure objective):约束树结构,使合并后的 chunk 语义更连贯,例如通过对比学习让相似句子在树中更早合并。 最终损失为两部分的加权和,模型同时学习高质量的节点 embedding 和更具语义凝聚力的树结构。

与同类方法的差异

与现有依赖外部 LLM 分块、固定上下文窗口扩展或层次化摘要的工作不同,SproutRAG 完全基于模型内部的注意力机制进行无损的多粒度索引和检索,既避免昂贵的 LLM 调用,又保留了原始文本的全部信息。

实验

实验设计

评估基準覆盖科学 (SCI-DOCS)、法律 (LegalBench-RAG) 和开放域问答 (HotpotQA, MS MARCO) 四个长文档 RAG 数据集。基线方法包括密集检索 (Dense X Retrieval)、元分块 (Meta-Chunking)、混合上下文 (MoC)、层次总结树 (RAPTOR) 以及 LightRAGPropRAG 等 9 种近期方案。主要指标为信息效率 (IE),衡量检索到的上下文对下游生成质量的提升效率。同时考察不同截断位置的检索召回率,并设计了消融实验,逐一移除树结构、分层波束搜索和联合训练,验证各模块贡献。

关键发现

  1. 全局最优:SproutRAG 在所有基准上平均 IE 提升 6.1%,尤其在需要多句证据的法律案例中增益显著。
  2. 无额外 LLM 成本:相比依赖 LLM 指导分块或总结的基线,SproutRAG 仅通过训练时学到的句子间注意力构建语义树,索引和检索阶段不调用 LLM。
  3. 多粒度受益:分层波束搜索允许同时召回细粒度句子粗粒度段落群,捕捉平层检索遗漏的跨句关联。消融实验证实,移除分层搜索后长文本问题性能下降最明显。
  4. 联合训练必要:分开优化嵌入和树结构会破坏语义对齐,联合目标同时提升稠密向量质量树结构合理性

对比解读

RAPTOR 的分层总结不同,SproutRAG 不生成压缩摘要,信息零损失;与 Meta-Chunking 的单层扩展不同,它提供可伸缩的多粒度上下文;与 LightRAG 等基于图谱的方法相比,注意力引导的树搜索更轻量且可端到端训练。这一设计思路对工程实践的启示:当文档语义层级可通过模型内部注意力信号挖掘时,便无需引入昂贵的外部逻辑,从而在成本、粒度与连贯性之间找到新平衡。

行业影响

落地场景

SproutRAG 的多粒度层次化检索天然适合处理长文档理解与问答,可嵌入以下产品线:

  • 企业知识库与法律 / 合规平台:对政策文件、合同、判例库进行句子级细粒度索引,支持精确条文引用与多证据聚合,不依赖外部 LLM 分块,显著降低检索延迟。
  • 学术文献与医疗文献助手:将科学论文按句子组织为二进制分块树,实现“摘要-段落-句子”多级检索,辅助研究人员快速定位关键论据,避免摘要压缩造成的信息丢失。
  • 垂直搜索与客服系统:在电商评论分析、金融研报解读等场景,通过对用户 query 的层次化 beam search,可同时返回高相关度的句子和其所在段落,提升答案的上下文连贯性。

商业价值

  • 降本:完全省去索引阶段 LLM 调用的成本,仅需少数注意力头的前向传播,大规模文档库的更新边际成本极低。
  • 增收与体验提升:多粒度检索提升答案的召回率和信息效率(IE),在需要精确证据链的场景(如法律咨询、医疗诊断辅助)可大幅降低误答风险,增强用户信任与产品粘性。
  • 效率指标:论文显示在四个基准上 IE 平均提升 6.1%,意味着在相同向量库规模下能提供更高质量的上下文,从而减少 RAG 管道后续的生成重试或人工校对成本。

与现有产品 / 工作流的接口

SproutRAG 可以作为“智能分块 + 层次化检索”的中间件插入现有 RAG 栈:

  • 索引阶段:替换传统固定长度分块或 LLM 分块,输出一棵二进制树结构,每个内部节点携带递进式嵌入。可与向量数据库(如 Milvus、Qdrant)集成,将多粒度嵌入存入同一集合的不同 partition。
  • 检索阶段:对外暴露类似 retrieve(query, k, beam_width) 的 API,返回 [chunk_embeddings, node_paths],下游生成器可直接拼接作为 LLM 上下文。支持动态粒度选择(granularity='sentence' | 'paragraph')。
  • 训练流程:端到端联合优化可直接融入现有模型微调管线,使用用户反馈信号(如点击、点赞)增强树结构与嵌入质量。无需修改语言模型本体,仅额外训练注意力头选择器和 chunking 策略网络。

具体落地 use case

  1. 法律科技平台(如 Robin AI、Legalese):在处理“根据某条款,承租人需在几个月内通知?”类问题时,传统分块可能切碎条款上下文。SproutRAG 通过注意力树将“通知期限”相关句子与前置条件句子编入同一上层节点,检索时可一次返回完整条款段,提高回答准确率并降低律师人工复核成本。

  2. 医疗多院长文档助手(如医院内部的 EHR 检索):病历、诊断报告通常篇幅长且互相关联稀疏。SproutRAG 的句子级索引可快速定位“患者曾用药物 X 出现副作用 Y”的散落描述,并向上回溯相关就诊段落,避免摘要模型丢失关键时序信息,辅助医生在随访中做出更安全的决策。

局限

  • - **训练依赖性**:SproutRAG 避免了检索时的 LLM 调用,但其端到端训练仍需大量标注查询-文档对,以学习注意力引导的树结构和嵌入。对于每个新领域或文档集,可能需要重新训练或微调,这限制了零样本迁移能力。相比基于 LLM 的通用分块方法,其适应性更窄,且联合优化的计算开销较高,在资源受限场景下难以快速部署。
  • - **注意力引导的鲁棒性**:框架核心依赖“学习哪些注意力头与层最能捕捉语义结构”,这一假设受基模型注意力质量影响较大。若文档结构复杂或注意力图存在噪声,构建的二叉分块树可能无法真实反映语义连贯性,导致检索性能下降。论文未在低质量注意力模型或非预训练语言模型上验证,方法对不同编码器架构的泛化性存疑。
  • - **评估范围与实用性**:实验仅在四个基准上评估,且重点报告信息效率(IE)提升 6.1%,未全面对比最新的层次化或图 RAG 方法(如 GraphRAG 在结构化信息上的表现)。层次化束搜索的检索延迟可能高于平面检索,但论文未深入讨论效率 trade-off。代码库 Star 数较低,尚缺乏大规模工业验证,实际可用性和长期维护情况不明。
论文Amirhossein Abaskohi2026-06-16原文

相关内容