SproutRAG：注意力引導的樹搜索與漸進式嵌入用於長文檔RAG

摘要

檢索增強生成（RAG）系統必須在檢索粒度與上下文連貫性之間取得平衡，現有方法透過LLM引導分塊、單層級上下文擴展或層次化摘要來應對此挑戰。這些方法在索引或檢索階段不同程度地依賴成本高昂的LLM調用，將上下文聚合限制在單一粒度層級，或透過摘要引入資訊損失。我們提出SproutRAG，一種注意力引導的層次化RAG框架，通過將句子級別的分塊組織成逐步增大但語義連貫的單元，並利用學習到的句子間注意力來構建二元分塊樹，從而解決此權衡問題。與依賴外部LLM、固定上下文擴展或有損摘要的先前方法不同，SproutRAG學習哪些注意力頭與層最能捕捉語義文檔結構，使得無需額外LLM調用或壓縮摘要即可實現多粒度檢索。在檢索時，SproutRAG使用層次化束搜索在多粒度下檢索候選項，捕捉超越平面檢索的多句子相關性。該框架通過聯合目標進行端到端訓練，同時改進嵌入與樹結構。在涵蓋科學、法律與開放域的四個基準實驗中，SproutRAG的資訊效率（IE）平均比最強基線提升6.1%。程式碼可於https://github.com/AmirAbaskohi/SproutRAG取得。

English

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.