SproutRAG: Aandacht-geleid boomzoeken met progressieve embeddings voor RAG met lange documenten

Samenvatting

Retrieval-augmented generation (RAG)-systemen moeten een balans vinden tussen retrievergranulariteit en contextuele coherentie, een uitdaging die bestaande methoden aanpakken via LLM-gestuurd chunking, contextuitbreiding op één niveau, of hiërarchische samenvatting. Deze benaderingen zijn op verschillende manieren afhankelijk van dure LLM-aanroepen tijdens het indexeren of ophalen, beperken contextaggregatie tot één granulariteitsniveau, of introduceren informatieverlies door samenvatting. Wij presenteren SproutRAG, een door aandacht gestuurd hiërarchisch RAG-raamwerk dat deze afweging aanpakt door zinsniveau-brokken te organiseren in progressief grotere maar semantisch coherente eenheden, waarbij gebruik wordt gemaakt van aangeleerde inter-zinsaandacht om een binaire chunking-boom te construeren. In tegenstelling tot eerdere benaderingen die afhankelijk zijn van externe LLM's, vaste contextuitbreiding of verliesgevende samenvatting, leert SproutRAG welke aandachtskoppen en -lagen de semantische documentstructuur het beste vastleggen, wat multi-granulariteitsretrieval mogelijk maakt zonder extra LLM-aanroepen of gecomprimeerde samenvattingen. Tijdens het ophalen gebruikt SproutRAG hiërarchisch beam search om kandidaten op meerdere granulariteitsniveaus te vinden, waarbij meerzinsrelevantie wordt vastgelegd die verder gaat dan vlakke retrieval. Het raamwerk wordt end-to-end getraind met een gezamenlijke doelstelling die zowel de embeddings als de boomstructuur verbetert. Experimenten op vier benchmarks in wetenschappelijke, juridische en open-domein omgevingen tonen aan dat SproutRAG de informatie-efficiëntie (IE) gemiddeld met 6,1% verbetert ten opzichte van de sterkste baseline. Code is beschikbaar op 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.