PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Amirhossein Abaskohi, Issam H. Laradji, Peter West, Giuseppe Carenini

arXiv:2606.18381v1cs.CLcs.IR

TL;DR

RAG must balance precise retrieval with contextual coherence, especially when relevant evidence spans multiple document sections. SproutRAG builds an attention-guided hierarchical tree for multi-granularity retrieval and improves information efficiency by 6.1% on average over the strongest baseline.

  • Problem

    RAG chunking must balance contextual coherence against precision, particularly when answering queries requiring evidence scattered across multiple document sections.

  • Method

    SproutRAG learns inter-sentence attention to build a binary tree with progressive embeddings and uses hierarchical beam search for multi-granularity retrieval.

  • Results

    6.1% average improvement in information efficiency over the strongest baseline was observed across four scientific, legal, and open-domain benchmarks.

  • Takeaways & Limitations

    SproutRAG offers a performance–efficiency tradeoff for multi-granularity retrieval while using fewer online tokens and lower latency than LLM-heavy systems.

  • Takeaways & Limitations

    The binary hierarchy may restrict grouping several sentences that jointly form a coherent semantic unit, and fixed offline trees may be less flexible for query-specific evidence.

Abstract

from arXiv · show

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.

1 Introduction

Long-document RAG must balance contextual coherence against retrieval precision because large chunks add noise while fine-grained chunks fragment semantics and cross-chunk relationships. SPROUTRAG addresses this with a learned attention-guided binary tree, progressive embeddings, and hierarchical beam search for multi-granularity retrieval without external LLM calls.

  • Long-document RAG segments documents because full-document inputs face context-length constraints and degraded attention over extended sequences.
  • Chunk granularity creates a trade-off: large chunks preserve coherence but add redundant noise, whereas fine-grained chunks improve precision but fragment semantics and interchunk relationships.
  • SPROUTRAG organizes sentence-level chunks into a learned attention-guided binary tree that preserves cross-chunk dependencies and supports retrieval at multiple semantic granularities.The tree is constructed bottom-up using SLLM representations and learned inter-sentence attention.
  • Progressive embeddings compositionally represent each subtree, while hierarchical beam search retrieves candidates across fine-grained leaves, mid-level nodes, and broader subtrees.
  • SPROUTRAG is trained end-to-end to jointly optimize retrieval quality and tree structure without external LLM calls during inference.

2 Related Work

Related work improves RAG by adapting chunk granularity, retrieval, or document structure, but existing methods retain flat retrieval units, use pairwise expansion, or rely on clustering and summaries. These approaches motivate hierarchical retrieval that preserves fine-grained information while supporting multiple levels.

  • Chunking and adaptive retrieval: Standard RAG commonly uses efficient fixed-length or delimiter-based chunking, but these rule-based splitters are insensitive to semantic boundaries.
  • Chunking and adaptive retrieval: ReflectiveRAG evaluates evidence sufficiency and reformulates queries, yet it leaves the underlying flat organization of retrieval units unchanged.
  • Chunking and adaptive retrieval: SAKI-RAG estimates inter-sentence attention with an SLLM and expands retrieved chunks with related sentences using pairwise expansion and LLM filtering.
  • Structured and hierarchical retrieval: RAPTOR recursively clusters chunks and summarizes clusters into a tree, enabling multi-level retrieval but potentially discarding fine-grained details through LLM-generated summaries.

3 SPROUTRAG

SproutRAG builds an attention-guided binary hierarchy over sentence-level chunks, using progressive embeddings to represent increasingly broad semantic units. It retrieves evidence through hierarchical beam search and jointly trains embeddings with the tree structure.

  • Indexing: SproutRAG encodes sentence-level chunks with a SLLM and learns head–layer attention weights to reduce proximity bias and capture semantic co-relevance.The weighted attention is symmetrized into mutual relation scores that guide hierarchy construction.
  • Indexing: The model builds the tree bottom-up by repeatedly merging the active node pair with the highest mutual attention score and computing progressive parent embeddings.Leaves preserve sentence-level precision, while internal nodes represent broader semantic units and retain long-range connections through single-linkage updates.
  • Retrieval: At retrieval, SproutRAG encodes the query and performs hierarchical beam search, expanding promising branches while retaining the top-b nodes at each depth.This focuses computation on relevant branches while allowing exploration of multiple semantically relevant document regions.
  • Retrieval: The candidate set combines visited leaves, internal nodes, and subtrees, enabling retrieval of either precise facts or broader multi-sentence context at query-appropriate granularities.Candidates are reranked by similarity or a lightweight reranker before the top-k chunks are passed to the answer generator.
  • Training: SproutRAG jointly optimizes contrastive retrieval embeddings and an attention regularizer that groups co-relevant sentence pairs into coherent, retrievable subtrees.After training, indexing uses the learned aggregation weights, while retrieval requires query encoding, tree traversal, and reranking without external LLM filtering or lossy summarization.

4 Experiments and Results

SproutRAG is evaluated across retrieval and end-to-end answer-generation benchmarks, where it improves information efficiency and precision while supporting broader evidence aggregation. Ablations show that joint supervision, learned attention-guided hierarchy, internal-node retrieval, beam search, and default hyperparameters are important to its performance.

  • Retrieval performance: SPROUTRAG achieves the highest IE and precision on all four retrieval benchmarks, improving IE over the strongest baseline by 8.06 points on Dragonball, 4.65 on SCI-DOCS, 4.90 on LegalBench-RAG, and 6.83 on MS MARCO.The results indicate broader supporting-context retrieval without the noise associated with overly large retrieval units.
  • Retrieval performance: SPROUTRAG preserves SAKI-RAG’s precision benefits while improving IE through a global tree that retrieves individual chunks, internal nodes, and subtrees.This avoids the limitations of pairwise expansion and enables multi-granularity evidence aggregation.
  • End-to-end evaluation: On HotpotQA, WebQuestions, and Dragonball, end-to-end evaluation compares answer quality and online efficiency, with PageIndex obtaining the highest final answer scores but requiring substantially more online computation.SPROUTRAG incurs upfront training, while its learned hierarchy generalizes across datasets without retraining and amortizes that cost rather than adding query-time overhead.
  • Ablation study: Removing retrieval or attention losses degrades retrieval, showing that query–evidence alignment and hierarchy shaping require complementary supervision.The not-trained variant performs worst across datasets, while removing retrieval loss substantially reduces recall and IE.
  • Ablation study: Uniform attention aggregation and embedding-similarity trees underperform, demonstrating that learned attention aggregation, internal-node retrieval, and beam search are crucial for balancing precision with contextual coverage.The ablations indicate that SLLM attention captures structural information beyond embedding similarity.
  • Hyperparameter sensitivity: The default beam width b = 5 and attention weight λ = 0.1 outperform tested alternatives overall: b = 3 lowers IE, b = 10 raises recall but lowers precision, and λ = 0.05 or λ = 0.20 underperform.The results favor a middle setting that balances explored semantic paths, precision, and tree-construction guidance.

5 Conclusion and Future Work

SPROUTRAG is an attention-guided hierarchical RAG framework that organizes sentence-level chunks into a learned tree for multi-granularity retrieval. It constructs a retrieval-oriented hierarchy with learned SLLM attention aggregation and uses hierarchical beam search over multiple node granularities.

  • Conclusion: SPROUTRAG organizes sentence-level chunks into a learned tree for multi-granularity retrieval.The framework is described as attention-guided and hierarchical.
  • Conclusion: Learned SLLM attention aggregation constructs a retrieval-oriented hierarchy instead of relying on fixed boundaries, context expansion, lossy summarization, or inference-time LLM filtering.The hierarchy is designed around retrieval rather than summarization.
  • Conclusion: Hierarchical beam search selects evidence from sentence leaves, internal nodes, and broader subtrees at inference time.This multi-level selection lets the retriever balance fine-grained and broader evidence.

Limitations

SproutRAG’s binary hierarchy may restrict grouping sentences that jointly form coherent semantic units, and the method requires upfront training for its SLLM and attention aggregation weights. The training cost is one-time and transfers across datasets.

  • Tree structure: The binary hierarchy may restrict grouping several jointly coherent sentences into one semantic unit.Multi-branch trees could better capture such many-to-many dependencies.
  • Training cost: SproutRAG requires upfront training for the SLLM and attention aggregation weights.This is a one-time cost that transfers across datasets.

A Benchmark and Baseline Details · A.1 Benchmark Dataset Details · LegalBench-RAG (Pipitone and Alami, 2024)

The evaluation spans retrieval-focused and end-to-end generation benchmarks covering scientific, legal, open-domain, multi-hop, factual, and multi-domain RAG settings. LegalBench-RAG specifically tests fine-grained retrieval of minimal legal evidence across diverse document types.

  • A Benchmark and Baseline Details: The study uses four retrieval-focused benchmarks: SCI-DOCS, LegalBench-RAG, Dragonball, and MS MARCO.These benchmarks cover scientific retrieval, legal retrieval, open-domain passage retrieval, and multi-domain RAG evaluation.
  • A.1 Benchmark Dataset Details: The end-to-end generation benchmarks are HotpotQA, WebQuestions, and Dragonball, spanning multi-hop reasoning, short-form factual QA, and multi-domain RAG.Together, the datasets cover scientific, legal, and open-domain retrieval as well as multiple generation settings.
  • A.1 Benchmark Dataset Details: SCI-DOCS tests scientific retrieval using dense, terminology-heavy abstracts containing multiple related concepts within short spans.It includes citation prediction, document classification, and recommendation tasks, and was introduced with SPECTER.
  • LegalBench-RAG (Pipitone and Alami, 2024): LegalBench-RAG contains 6,858 query-answer pairs over legal documents including NDAs, M&A agreements, commercial contracts, and privacy policies.It is designed for legal RAG retrieval and emphasizes precise snippet retrieval rather than generally relevant document retrieval.
  • LegalBench-RAG (Pipitone and Alami, 2024): LegalBench-RAG requires identifying minimal legal evidence, making it a strong test of fine-grained precision.Its focus differs from broad document retrieval by evaluating precise snippet selection.
  • A.1 Benchmark Dataset Details: Dragonball evaluates retrieval and generation across finance, legal, and medical scenarios in English and Chinese.Its heterogeneous domains, long evidence contexts, and domain-specific terminology test relevant-evidence recovery without excessive distractor context.
  • A.1 Benchmark Dataset Details: HotpotQA requires reasoning over multiple supporting documents or facts, while WebQuestions evaluates short-entity or short-phrase answers with token-level F1.HotpotQA provides sentence-level supporting-fact annotations; WebQuestions tests precise retrieval without unnecessary context.

A.2 Baseline Details

SPROUTRAG is compared with retrieval-oriented and system-level RAG baselines using original-paper configurations and shared QWEN3 generator and reranker settings where applicable. The baselines span proposition retrieval, adaptive chunking, hierarchical and graph retrieval, sentence-level attention expansion, and end-to-end adaptive or efficiency-focused systems.

  • Evaluation setup: Baselines comprise efficient retrieval methods for retrieval evaluation and system-level RAG methods for end-to-end generation comparison.Unless otherwise stated, configurations follow the original papers; methods requiring generation or reranking use the same QWEN3-8B generator and QWEN3-RERANKER-4B reranker.
  • Retrieval-oriented baselines: Dense X Retrieval, Meta-Chunking, and MoC improve atomicity, boundary coherence, or adaptive segmentation but do not explicitly model cross-chunk dependencies or retrieval-time sentence hierarchies.Meta-Chunking is evaluated with both perplexity-based and margin-sampling-based boundary decisions.
  • Structured retrieval baselines: RAPTOR, LightRAG, and PropRAG provide hierarchical or structured retrieval through summaries, graphs, or proposition paths, differing from SPROUTRAG’s attention-guided sentence hierarchy.RAPTOR’s summaries can introduce information loss and indexing cost, while PropRAG uses beam traversal over a proposition graph.
  • Attention-based baseline: SAKI-RAG is the closest baseline because it uses sentence-level attention, but it performs pairwise expansion with retrieval-time LLM filtering instead of global multi-granularity tree retrieval.SPROUTRAG converts learned attention into a global tree and avoids inference-time LLM filtering.
  • System-level baselines: GraphRAG, ReflectiveRAG, PageIndex, and REFRAG are included in end-to-end comparisons for graph reasoning, adaptive evidence use, LLM-guided search, or decoding-time context efficiency.PageIndex is excluded from efficient retrieval evaluation because its online search and evidence construction are LLM-heavy; REFRAG focuses on compressing and selectively expanding retrieved context during decoding.

B Retrieval Performance at Different Cutoffs

Across six benchmarks and cutoffs k=1, 3, and 5, precision and recall increase monotonically with k, while SPROUTRAG consistently outperforms all non-oracle baselines at every depth.

  • Cutoff Evaluation: Tables 4–6 evaluate retrieval performance at k=1, k=3, and k=5 across all six benchmarks.The three tables correspond respectively to cutoffs k=1, k=3, and k=5.
  • Cutoff Evaluation: Precision and recall increase monotonically with k for every method because retrieving more documents expands relevant-passage coverage.This trend is reported across the evaluated cutoffs and methods.
  • Method Comparison: SPROUTRAG outperforms all non-oracle baselines at every retrieval depth, preserving the same method ordering across cutoffs.The passage attributes this consistency to gains from topic-guided retrieval.

C Qualitative Analysis: Recovering Multi-Sentence Legal Evidence

SproutRAG’s internal tree-node retrieval recovers multi-sentence legal evidence that pairwise or local chunking misses. In the limitation-of-liability example, it retrieves a clause-level unit containing four relevant sentences, enabling a legally complete answer.

  • Multi-Sentence Evidence: SAKI-RAG links the damage-exclusion and liability-cap sentences but misses the later exception sentence.The omitted exception is essential for a legally complete answer.
  • Multi-Sentence Evidence: SPROUTRAG retrieves internal node v1:4, grouping all four relevant sentences into one clause-level unit.The unit supports answering both the liability-limit and exception parts of the query.
  • Legal Completeness: SPROUTRAG enables an answer covering excluded damages, a monetary liability cap, legal-theory scope, and specified carve-outs.The retrieved node captures the evidence needed to state that the limitation does not apply to the carve-outs.
  • Qualitative Comparison: MoC retrieves a locally coherent but incomplete chunk for the limitation-of-liability query.The qualitative comparison contrasts local coherence with the broader evidence recovered by SPROUTRAG.
Loading 2606.18381v1…