← 返回信息流
技术博客arXiv cs.CL·1 小时前

SproutRAG：基于注意力引导树搜索的长文档RAG新方法

原标题：SproutRAG: Attention-Guided Tree Search with Progressive Embeddings for Long-Document RAG

速览

SproutRAG是一种注意力引导的分层RAG框架，通过将句子级块组织成语义连贯的单元，构建二叉块树以捕捉文档结构。该方法无需外部大模型调用或压缩摘要，即可实现多粒度检索，有效平衡检索粒度与上下文连贯性。在科学、法律和开放域基准测试中，其信息效率平均比最强基线提升6.1%。

AI 深度解读

SproutRAG：基于注意力引导的树搜索与渐进式嵌入，解决长文档 RAG 难题

背景

检索增强生成（Retrieval-Augmented Generation, RAG）系统在面对长文档处理时，始终面临着一个核心矛盾：检索粒度（Retrieval Granularity）与上下文连贯性（Contextual Coherence）之间的平衡。

现有的主流解决方案通常采取以下几种策略，但各自存在明显的局限性：

  1. LLM 引导的分块（Chunking）：依赖大型语言模型在索引或检索阶段进行文档切分，导致高昂的计算成本。
  2. 单层上下文扩展：仅能在单一的粒度层级上聚合上下文，难以捕捉文档中多层次的语义结构。
  3. 层级摘要（Hierarchical Summarization）：通过生成摘要来压缩信息，但这往往伴随着信息丢失（Information Loss），导致检索精度下降。

这些方法要么过度依赖外部 LLM 调用，要么限制了上下文聚合的灵活性，要么引入了不可逆的信息损耗。为了解决这一权衡难题，研究人员提出了 SproutRAG，一种无需额外 LLM 调用、能实现多粒度检索的新型框架。

核心内容

SproutRAG 是一种注意力引导的分层 RAG 框架，其核心创新在于利用学习到的句间注意力机制，将句子级别的块（Sentence-level chunks）组织成逐渐增大但语义连贯的单元。

1. 构建二分块树（Binary Chunking Tree）

不同于传统方法依赖外部 LLM 或固定规则，SproutRAG 通过以下方式构建文档结构：

  • 学习注意力机制：模型自动学习哪些注意力头（Attention Heads）和层（Layers）最能捕捉文档的语义结构。
  • 动态树构建：基于上述注意力权重，构建一个二分的块树结构。这种结构允许文档从细粒度的句子逐步聚合为更粗粒度的语义单元，形成一种层次化的表示。

2. 渐进式嵌入与多粒度检索

  • 渐进式嵌入（Progressive Embeddings）：在树的不同层级，嵌入表示逐渐涵盖更大的语义范围，但保持连贯性。
  • 分层束搜索（Hierarchical Beam Search）：在检索阶段，SproutRAG 使用分层束搜索算法，在多个粒度层级上同时检索候选项。这使得系统能够捕捉到超越扁平化检索（Flat Retrieval）的多句子相关性，从而更全面地理解用户查询与文档片段之间的关系。

3. 端到端联合训练

SproutRAG 采用联合目标函数进行端到端训练，该目标函数同时优化：

  • 嵌入质量：确保不同粒度的嵌入能准确反映语义。
  • 树结构：确保二分树的构建符合文档的内在语义逻辑。

4. 无需额外 LLM 调用

与依赖外部 LLM 进行分块或摘要的方法不同，SproutRAG 完全基于内部学习到的参数和注意力机制工作。这意味着在索引和检索阶段，无需额外的 LLM 调用，也无需生成压缩摘要，从而避免了高昂的成本和信息丢失。

关键要点

  • 解决粒度与连贯性的权衡：SproutRAG 通过构建分层结构，同时兼顾了细粒度的精确检索和粗粒度的上下文连贯性。
  • 基于注意力的结构学习：模型并非使用固定规则或外部 LLM，而是通过学习内部注意力机制来自动发现文档的最佳分割点和层级结构。
  • 二分块树结构：将句子级块组织成二叉树形式，每个节点代表一个语义连贯的单元，从叶子节点（句子）到根节点（文档整体）粒度逐渐增加。
  • 分层束搜索检索：在检索时，算法在树的不同层级并行搜索，能够捕获跨越多句子的复杂语义相关性，优于传统的扁平检索。
  • 端到端联合优化：嵌入模型和树结构通过单一联合目标函数共同训练，确保两者相互促进，达到最优性能。
  • 零额外 LLM 开销：整个流程（索引、检索）不依赖外部 LLM 调用，显著降低了延迟和成本，同时避免了摘要带来的信息损失。
  • 显著的性能提升：在涵盖科学、法律和开放域设置的四个基准测试中，SproutRAG 平均将信息效率（Information Efficiency, IE）提升了 6.1%，优于最强的基线模型。

意义与影响

SproutRAG 的提出标志着长文档 RAG 系统从“依赖外部重型模型”向“内部结构化感知”的重要转变。

  1. 效率与成本的优化：通过消除对索引和检索阶段外部 LLM 调用的依赖，SproutRAG 大幅降低了计算资源消耗和响应延迟，使得长文档 RAG 在大规模生产环境中更具可行性。
  2. 信息保真度的提升：避免了摘要生成过程中的信息压缩和丢失，确保了检索内容的完整性和准确性，这对于法律、医疗等对信息精度要求极高的领域尤为重要。
  3. 语义理解的深化：通过引入分层树结构和注意力引导，模型能够更深刻地理解文档的内在逻辑和语义层次，从而在复杂查询下提供更精准、更连贯的上下文支持。
  4. 技术范式的创新：SproutRAG 证明了通过优化嵌入结构和检索算法，可以在不增加模型复杂度的前提下显著提升 RAG 性能，为后续研究提供了新的思路，即如何通过内部机制而非外部工具来解决长文本处理的难题。

随着长上下文窗口成为 LLM 的标准配置，如何高效、精准地利用这些长文档仍是关键挑战。SproutRAG 提供了一种高效、低成本且高保真的解决方案，有望成为未来长文档 RAG 系统的重要基准。

相关推荐

查看原文 →arxiv.org