SproutRAG:基于渐进嵌入的注意力引导树搜索技术,适用于长文档RAG场景
SproutRAG: Attention-Guided Tree Search with Progressive Embeddings for Long-Document RAG
摘要
检索增强生成系统需要在检索的精细度与上下文一致性之间取得平衡。现有方法通过基于大语言模型的分块策略、单层上下文扩展或分层总结等方式来应对这一挑战。这些方法在索引或检索过程中需要依赖昂贵的大语言模型调用,将上下文聚合限制在单一精细度级别,或者通过总结过程导致信息损失。我们提出了SproutRAG这一基于注意力机制的分层检索增强框架,它通过将句子级片段组织成更大但语义上连贯的单位来解决这一矛盾问题,同时利用句子间注意力机制构建二元分块结构。与那些依赖外部大语言模型、固定上下文扩展或有损总结的方法不同,SproutRAG能够学习到哪些注意力头和层次最有利于捕捉文档的语义结构,从而实现多精细度检索,而无需额外的大语言模型调用或压缩后的总结结果。在检索过程中,SproutRAG使用分层束搜索算法来检索多个精细度的候选结果,从而能够捕捉到超出简单检索范围的多句相关性信息。该框架采用端到端训练方式,其目标就是同时提升嵌入向量和树结构的质量。在四个不同领域的测试基准上的实验表明,SproutRAG平均能将信息效率提高6.1%。代码可访问于https://github.com/AmirAbaskohi/SproutRAG。
English Abstract
Retrieval-augmented generation (RAG) systems must balance retrieval granularity with contextual coherence, a challenge that existing methods address through LLM-guided chunking, single-level context expansion, or hierarchical summarization. These approaches variously depend on costly LLM calls during indexing or retrieval, limit context aggregation to a single granularity level, or introduce information loss through summarization. We present SproutRAG, an attention-guided hierarchical RAG framework that addresses this trade-off by organizing sentence-level chunks into progressively larger but semantically coherent units, using learned inter-sentence attention to construct a binary chunking tree. Unlike prior approaches that rely on external LLMs, fixed context expansion, or lossy summarization, SproutRAG learns which attention heads and layers best capture semantic document structure, enabling multi-granularity retrieval without additional LLM calls or compressed summaries. At retrieval time, SproutRAG uses hierarchical beam search to retrieve candidates at multiple granularities, capturing multi-sentence relevance beyond flat retrieval. The framework is trained end-to-end with a joint objective that improves both embeddings and tree structure. Experiments across four benchmarks spanning scientific, legal, and open-domain settings demonstrate that SproutRAG improves information efficiency (IE) by 6.1% on average over the strongest baseline. Code is available on https://github.com/AmirAbaskohi/SproutRAG.