SproutRAG: Attention-Guided Tree Search with Progressive Embeddings for Long-Document RAG

2026-06-16

AI Summary

This work addresses the challenge in long-document retrieval-augmented generation (RAG) of balancing retrieval granularity with contextual coherence. The authors propose an attention-guided hierarchical RAG framework that dynamically constructs a binary chunking tree using internal Transformer attention mechanisms, enabling the organization of multi-granular yet semantically coherent units without relying on external large models or lossy summarization. The approach jointly optimizes embedding representations and tree structure in an end-to-end manner and employs hierarchical beam search to capture relevance across multiple granularities. Evaluated on four benchmarks spanning scientific, legal, and open-domain settings, the method achieves a 6.1% average improvement in information efficiency (IE) over the strongest baseline.

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.

Problem

• Retrieval-Augmented Generation
• Long-Document RAG
• Contextual Coherence
• Retrieval Granularity
• Information Loss

Innovation

• attention-guided tree search
• progressive embeddings
• multi-granularity retrieval
• hierarchical RAG
• end-to-end training