Amirhossein Abaskohi, Issam H. Laradji, Peter West, Giuseppe Carenini
6 min
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.
Alex: And once it has that tree structure built, how does it use it to answer a question?
Sam: It starts at the top of the tree and works its way down. At each level, it keeps the most promising branches and discards the rest. Think of it like navigating a library where you first pick the right floor, then the right shelf, then the right book — rather than searching every shelf on every floor at once. The system moves from broad to specific, and because the map was built in advance, it doesn't need to do heavy calculations while you're waiting for an answer.
Alex: That's what makes it faster, then. The expensive work happens beforehand.
Sam: Exactly. And that's also what makes it more reliable. The model is trained with two specific goals. The first is to pull questions and relevant documents closer together in its internal sense of meaning, while pushing unrelated things apart. It's a bit like training a search dog — you show it the scent you want, and you also show it things that definitely aren't that scent, so it learns the difference precisely.
Alex: And the second goal?
Sam: The second is what the researchers call a structure objective. The model is specifically rewarded when it assigns high attention to sentences that are genuinely related to each other. This isn't just about finding the right answer — it's about building a tree where the branches are organized by topic in a meaningful way. The training pushes the model to make its internal map accurate, not just good enough.
Alex: So when someone asks a question, the system isn't re-reading the whole document. It's walking down the paths in the map that are most likely to lead to the answer.
Sam: That's it. And because the map captures both fine-grained details and broader context, the system can match the level of detail to what the question actually needs. A question about a specific clause in a legal contract gets a precise, narrow answer. A question about the general argument of a scientific paper gets a broader response.
Alex: Does the evidence suggest this actually works better than existing approaches?
Sam: The study reports a meaningful improvement. The researchers measured what they call information efficiency — essentially a score of how much relevant evidence the system finds relative to the computational effort it takes. Across tests on scientific and legal documents, SPROUTRAG improved that score by over six percent compared to the strongest existing methods. And the researchers describe that as a consistent pattern across their tests, not a single outlier result.
Alex: Six percent might not sound dramatic, but in a domain like legal contracts or medical research, where a single missing sentence can change the meaning of everything, that kind of reliability matters.
Sam: It does. And the efficiency gain is meaningful too. By building the map offline and navigating it cleanly, the system avoids the extra, expensive AI calls that slow down many current approaches. It's not just more accurate — it's faster and cheaper to run.
Alex: So the core insight is really about giving the AI a better map to follow, rather than making it re-read everything from scratch every time.
Sam: That's the clearest way to put it. Most of the difficulty in large document search comes from not knowing where to look. By teaching the model to understand the structure of a document — not just its content — SPROUTRAG gives the search process a genuine starting point. It's a more precise way to handle the scale of information that real-world AI systems are expected to work with.
Alex: It's a good reminder that a lot of progress in AI isn't always about building bigger models. Sometimes it's about giving existing ones better tools to organize what they already know.
Sam: Well put. And for anyone curious to go deeper, the full paper is available if you want to see how the training and evaluation were structured in detail.
Alex: Thanks for listening to ResearchPod.