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

# SproutRAG：基于注意力引导的树搜索与渐进式嵌入的长文档RAG
来源：https://arxiv.org/html/2606.18381
Amirhossein Abaskohi¹, Issam H. Laradji¹,², Peter West¹, Giuseppe Carenini¹  
¹不列颠哥伦比亚大学，²ServiceNow Research

###### 摘要

检索增强生成（RAG）系统必须在检索粒度与上下文连贯性之间取得平衡，现有方法通过LLM引导的分块、单层上下文扩展或分层摘要来解决这一挑战。这些方法在索引或检索时依赖昂贵的LLM调用，将上下文聚合限制在单一粒度级别，或通过摘要丢失信息。本文提出 **SproutRAG**，一种注意力引导的分层RAG框架，通过将句子级块组织为逐步增大但语义连贯的单元来解决这一权衡问题，利用学习到的句子间注意力构建二叉分块树。与依赖外部LLM、固定上下文扩展或有损摘要的先前方法不同，SproutRAG学习哪些注意力头与层最能捕捉语义文档结构，从而在不增加LLM调用或压缩摘要的情况下实现多粒度检索。在检索时，SproutRAG使用分层波束搜索在多个粒度上检索候选，捕获超越平面检索的多句子相关性。该框架通过联合目标进行端到端训练，同时改进嵌入与树结构。在涵盖科学、法律和开放域四个基准上的实验表明，SproutRAG在信息效率（IE）上平均比最强基线提升6.1%¹。  
¹代码可在GitHub (https://github.com/AmirAbaskohi/SproutRAG) 获取。

![[无标题图片]](https://arxiv.org/html/2606.18381v1/figures/logo.png)  
**SproutRAG: 基于注意力引导的树搜索与渐进式嵌入的长文档RAG**

## 1 引言

检索增强生成（RAG）已成为将大型语言模型（LLM）锚定于外部知识的主流范式，有助于减少幻觉、支持领域特定推理，并提升知识密集型任务的表现（Lewis等，2020；Augenstein等，2024）。随着LLM越来越多地应用于涉及长文档的复杂任务（Jin等，2025b），由于上下文长度限制和长序列中注意力退化（Jin等，2025a；Liu等，2024），直接将整个文档作为输入变得不切实际。因此，RAG框架将文档分割成块，并检索最相关的片段，构建聚焦且高质量的生成证据。

检索步骤的有效性关键取决于文档如何分割。大块保留上下文连贯性，但引入冗余噪声稀释关键信息；而细粒度块提供精确性，但遭受语义碎片化和块间关系断裂（Tao等，2025；Zhao等，2025a）。这一问题在跨段落检索中尤为突出，因为回答查询需要综合多个文档部分的信息，如多跳推理和摘要任务（Liu等，2025）。

近期研究从多个方向解决这一挑战。SAKI-RAG（Tao等，2025）使用小型语言模型（SLLM）合并语义相关的句子对，并依赖外部LLM过滤检索候选。然而，将这种成对扩展扩展到多块相关性会大大增加候选空间，使LLM过滤变得昂贵。LLM引导的分块方法如Meta-Chunking（Zhao等，2025b）和MoC（Zhao等，2025a）改善了分割质量，但在分块后丢弃了跨块依赖。分层方法如RAPTOR（Sarthi等，2024）通过聚类和摘要支持多粒度检索，但聚类将组内块视为可互换，摘要可能丢失证据。基于图的方法如GraphRAG（Edge等，2025）建模实体关系，但当细粒度块包含稀疏实体信息时效果较差。

![图1](https://arxiv.org/html/2606.18381#S1.F1)  
图1: SproutRAG将长文档分割为句子级块，利用SLLM注意力识别语义相关句子，并将其组织成注意力引导的二叉树。检索时从细粒度叶子、中层节点和更大子树中选择证据，在保持精度的同时恢复连贯性。

本文提出 **SproutRAG**，一个注意力引导的分层RAG框架，将句子级块组织成学习到的文档结构，保留跨块依赖，同时避免检索期间的LLM推理开销。如图1所示，该结构支持在多个语义粒度上进行检索。SproutRAG使用SLLM以句子粒度编码文档，并自底向上构建二叉树，合并顺序由跨Transformer头和层的句子间注意力的学习加权聚合决定。这种聚合取代了朴素的均匀平均（我们发现均匀平均会引入邻近性偏差，削弱全局树索引）；相反，可学习的标量权重允许模型发现哪些头类型最能反映文档结构的语义相关性。每个内部节点存储一个渐进式嵌入，组合表示其子树，通过分层波束搜索在所有树层级收集候选，实现多粒度检索。SproutRAG经过端到端训练，联合优化检索质量和树结构，在任何阶段都不需要外部LLM调用。它能够捕获成对或平面检索方法无法捕捉的涌现性多句子相关性，同时保持足够高效以应用于部署。我们在四个覆盖科学、法律和开放域基准上评估SproutRAG。平均而言，SproutRAG在信息效率（IE）上比最强基线提升6.1%，尤其是在证据常常分散在段落间的情况下。

总之，本文的贡献如下：  
(1) 提出SproutRAG，一个注意力引导的分层RAG框架，利用学习到的句子间注意力在句子级块上构建二叉树，无需在推理时进行任何LLM调用即可实现多粒度检索。  
(2) 识别并解决句子级Transformer中均匀注意力平均引入的邻近性偏差，将其替换为学习权重聚合，使模型能够发现哪些注意力头最能反映文档结构的语义相关性。  
(3) 引入联合训练目标，同时改进检索质量和树结构，消除在流水线任何阶段使用外部LLM过滤或有损摘要的需要。

## 2 相关工作

#### 分块与自适应检索

RAG的有效性强烈依赖于文档如何分割为可检索单元。标准RAG流水线通常依赖基于规则的分割器，如固定长度或基于分隔符的分块，这些方法高效但对语义边界不敏感（Team，2024）。近期方法旨在改善这种粒度选择。Late-Chunking（Günther等，2025）在形成块嵌入之前对令牌表示进行上下文化，而Meta-Chunking（Zhao等，2025b）和MoC（Zhao等，2025a）使用LLM信号或路由机制产生更自适应的分块边界。Dense X Retrieval（Chen等，2024）通过将文本分解为原子命题向更细粒度迈进，提高了精确性但削弱了更广泛的上下文连续性。其他方法在块形成后调整检索。ReflectiveRAG（Verma等，2026）引入了自反思检索循环，评估证据充分性并重新表述查询以改善事实根基，但它没有改变检索单元的底层平面组织。与SproutRAG最相关的是SAKI-RAG（Tao等，2025），它使用SLLM估计句子间注意力并扩展检索到的块相关句子。与SAKI-RAG的成对扩展加LLM过滤不同，SproutRAG构建了一个全局句子级层次结构，支持多粒度检索，无需推理时LLM调用。

#### 结构化与分层检索

除了平面块检索，结构化RAG方法将文档内容组织成更高层次的表示。RAPTOR（Sarthi等，2024）递归聚类块并总结每个聚类为一棵树，支持从多个层次检索；然而，其结构基于嵌入空间聚类并依赖LLM生成的摘要，可能丢弃细粒度细节。基于图的方法如GraphRAG（Edge等，2025）和LightRAG（Guo等，2025）通过实体和关系表示文档，支持基于遍历的检索，但依赖于成功的实体提取和关系构建。PropRAG（Wang和Han，2025）用命题替换实体三元组，并在命题路径上执行无LLM的波束搜索，而Beam Retrieval（Zhang等，2024）展示了为多跳段落检索维护多个检索假设的好处。PageIndex（Zhang等，2025a）同样探索了基于推理的、无向量的文档树结构检索，但依赖于文档级结构组织而非学习到的句子级注意力。SproutRAG则在句子级块上构建一个注意力引导的二叉树，具有组合内部节点和所有节点的联合检索。这在不使用有损摘要、以实体为中心的结构或外部LLM调用的情况下保留了跨句子依赖。

## 3 SproutRAG

如图2所示，SproutRAG将平面块检索替换为基于句子级块的 **训练注意力引导层级结构**。在离线索引阶段，SLLM编码文档，提供句子嵌入和句子间注意力信号。这些信号通过可学习的头-层权重聚合，并用于构建二叉树，其中叶子表示细粒度块，内部节点存储合并句子组的渐进式嵌入。在线检索阶段，SproutRAG编码查询并执行分层波束搜索，从叶子、内部节点和子树收集候选。如第3.3节所述，该框架通过联合目标进行训练，同时改进检索质量和树结构，无需在检索期间使用外部LLM调用。

![图2](https://arxiv.org/html/2606.18381#S3.F2)  
图2: SproutRAG概览。在离线索引阶段（阶段1），文档被分割为句子级块，并用SLLM编码获得句子嵌入和句子间注意力。通过对注意力头和层的学习聚合引导自底向上树构建，产生一个注意力树，叶子为句子嵌入，内部节点为渐进式嵌入。在线检索阶段（阶段2），编码查询并通过分层波束搜索遍历树，从多个层级收集候选，然后通过相似性重排序选择前k个块用于答案生成。

### 3.1 注意力引导索引

给定文档 D，我们首先将其分割为句子级块 S = {s₁, ..., sₙ}。我们使用SLLM编码整个序列，获得上下文化的句子嵌入 {e(sᵢ)}ⁿᵢ₌₁以及所有层和头的注意力矩阵。对于层 l 和头 h，记对应的注意力矩阵为 Attn⁽ˡ·ʰ⁾ ∈ ℝⁿ×ⁿ。

对所有头和层进行均匀平均会过度强调局部句子邻近性，因为某些注意力头主要捕获序列模式（Voita等，2019）。为减少这种邻近性偏差，SproutRAG学习跨头和层的权重聚合：

Aᵢⱼ = Σₗ₌₁ᴸ Σₕ₌₁ᴴ wₗₕ Attn⁽ˡ·ʰ⁾ᵢⱼ, (1)

其中 wₗₕ 定义为：

wₗₕ = exp(αₗₕ) / Σₗ‘₌₁ᴸ Σₕ’₌₁ᴴ exp(αₗ‘ₕ’), (2)

可学习标量 αₗₕ 允许模型强调更好地捕获语义相关性的注意力头。然后我们对称化聚合注意力以获得相互关联分数：

Mᵢⱼ = (Aᵢⱼ + Aⱼᵢ) / 2. (3)

树采用自底向上构建。初始时，每个句子块是一个叶子节点。每一步，我们合并具有最高相互注意力分数的活跃节点对。父节点的嵌入计算为其子节点的渐进式嵌入：

e(p) = (e(u) + e(v)) / 2, (4)

其中 u 和 v 是被合并的子节点。合并后，新父节点继承其与每个剩余节点的最强关系：

Mₚᵣ = max(Mᵤᵣ, Mᵥᵣ). (5)

这种单链接更新在层级增长时保留长程语义连接。结果是一个注意力树 T，其叶子保留句子级精度，内部节点表示更广泛的语义单元。

### 3.2 分层检索

给定查询 q，SproutRAG首先使用与索引阶段相同的SLLM编码查询，获得查询嵌入 e(q)。然后在注意力引导的二叉树上进行检索，每个节点 v 表示特定粒度的文档片段。叶子节点对应句子级块，而内部节点表示逐步增大的语义相关句子组。这使得SproutRAG能够以最适用于查询的级别检索证据，而不是依赖单一的固定粒度。