Technical Summary: SPROUTRAG

Problem Statement

Retrieval-augmented generation (RAG) systems face a fundamental trade-off between retrieval granularity and contextual coherence. Existing methods struggle to balance these competing needs:

• Large chunks preserve context but introduce redundant noise that dilutes key information.
• Fine-grained chunks offer precision but suffer from semantic fragmentation and broken inter-chunk relationships, particularly in cross-paragraph retrieval tasks (e.g., multi-hop reasoning).
• Current solutions often rely on costly Large Language Model (LLM) calls during indexing or retrieval (e.g., LLM-guided chunking, pairwise expansion with LLM filtering), limit context aggregation to a single granularity level, or introduce information loss through hierarchical summarization (e.g., RAPTOR).

The paper argues that there is a need for a framework that can organize document structure to support multi-granularity retrieval without incurring the inference-time overhead of external LLMs or the information loss of summarization.

Methodology: SPROUTRAG

SPROUTRAG (Attention-Guided Tree Search with Progressive Embeddings) is a hierarchical RAG framework that constructs a learned binary tree over sentence-level chunks. It operates in two phases:

1. Attention-Guided Indexing (Offline)

• Sentence-Level Encoding: A Sentence-Level Large Language Model (SLLM) encodes the document to produce contextualized sentence embeddings and inter-sentence attention matrices across all transformer layers and heads.
• Learned Attention Aggregation: To address the "proximity bias" inherent in uniform averaging of attention heads (which overemphasizes sequential adjacency), SPROUTRAG learns a weighted aggregation of attention heads and layers. Learnable scalar weights ($\alpha_{l,h}$) allow the model to identify which specific attention heads best capture semantic co-relevance for document structure.
• Tree Construction: A binary tree is built bottom-up. Initially, each sentence is a leaf. At each step, the pair of active nodes with the highest mutual attention score is merged.
• Progressive Embeddings: Internal nodes store "progressive embeddings" computed as the average of their children's embeddings ($e_p = (e_u + e_v)/2$). This allows internal nodes to compositionally represent broader semantic units.
• Single-Linkage Update: When merging nodes, the parent inherits the strongest relation to remaining nodes, preserving long-range semantic connections as the hierarchy grows.

2. Hierarchical Retrieval (Online)

• Query Encoding: The query is encoded using the same SLLM.
• Hierarchical Beam Search: Retrieval traverses the tree using a beam search strategy. Starting from the root, the algorithm expands children of the current beam nodes and retains the top-$b$ most relevant nodes based on cosine similarity.
• Multi-Granularity Collection: Unlike flat retrieval, SPROUTRAG collects candidates from all levels of the tree traversal (leaves, internal nodes, and subtrees) that exceed a learned similarity threshold.
• Reranking: The collected candidates are reranked, and the top-$k$ chunks are passed to the generator. This allows the system to retrieve precise facts (leaves) or broader multi-sentence context (internal nodes) as appropriate for the query.

Joint Training Objective

The framework is trained end-to-end with a joint loss function:

1. Retrieval Objective ($L_{ret}$): A contrastive learning loss that aligns query embeddings with positive passages and separates them from hard negatives.
2. Structure Objective ($L_{attn}$): An attention regularizer that encourages co-relevant sentence pairs (those jointly supporting a query) to receive high mutual attention, directly shaping the induced tree structure.

Key Contributions

1. SPROUTRAG Framework: An attention-guided hierarchical RAG system that constructs a binary tree over sentence-level chunks using learned inter-sentence attention, enabling multi-granularity retrieval without any LLM calls at inference time.
2. Addressing Proximity Bias: The identification and mitigation of proximity bias in sentence-level transformers. The paper replaces naive uniform attention averaging with a learned weighted aggregation, allowing the model to discover which attention heads best reflect semantic co-relevance for document structure.
3. Joint Training: A novel joint training objective that simultaneously optimizes retrieval quality and tree structure, eliminating the need for external LLM filtering or lossy summarization at any pipeline stage.

Experimental Results

The authors evaluated SPROUTRAG on four benchmarks spanning scientific (SCI-DOCS), legal (LegalBench-RAG), open-domain (MS MARCO), and multi-domain (Dragonball) settings, as well as end-to-end generation tasks (HotpotQA, WebQuestions).

• Retrieval Performance: SPROUTRAG achieved the highest Information Efficiency (IE) (Recall $\times$ Precision) across all datasets, improving by an average of 6.1% over the strongest baseline. Specifically, it improved IE by 8.06 points on Dragonball, 4.65 on SCI-DOCS, 4.90 on LegalBench-RAG, and 6.83 on MS MARCO.
• Precision vs. Recall: Unlike methods that sacrifice precision for recall (or vice versa), SPROUTRAG achieved the best precision on every benchmark while maintaining high recall. This suggests the hierarchy effectively retrieves broader supporting context without introducing noise.
• Efficiency: In end-to-end generation tasks, SPROUTRAG offered the strongest performance-efficiency trade-off. It outperformed heavy reasoning-based systems (like GraphRAG and PageIndex) while using significantly fewer online tokens per query (4.38K) and lower latency (193 ms).
• Ablation Studies:
  • Removing the joint training objectives significantly degraded performance, confirming the necessity of both retrieval alignment and structure regularization.
  • Replacing learned attention aggregation with uniform averaging or embedding-similarity trees resulted in lower performance, validating the importance of learned attention signals for tree construction.
  • Leaf-only retrieval maintained high precision but suffered in recall and IE, demonstrating the value of retrieving internal nodes.

Significance and Claims

The paper claims that SPROUTRAG successfully addresses the trade-off between retrieval granularity and contextual coherence. By organizing sentence-level chunks into a learned, attention-guided hierarchy, the framework enables multi-granularity retrieval that captures emergent multi-sentence relevance which flat or pairwise methods miss.

Crucially, the authors emphasize that this is achieved without inference-time LLM calls or lossy summarization. The framework generalizes across diverse domains (scientific, legal, open) after a one-time training stage, amortizing the training cost. The results suggest that learned attention signals can effectively model document structure for retrieval, offering a path toward efficient, high-quality RAG systems that rival more computationally expensive, LLM-heavy approaches.

The paper acknowledges limitations, noting that the current binary tree structure may be restrictive for many-to-many dependencies and that the upfront training cost is higher than using off-the-shelf retrievers without adaptation. However, it positions SPROUTRAG as a significant step toward efficient, structure-aware retrieval.