SproutRAG: Aufmerksamkeitsgesteuerte Baumsuche mit progressiven Einbettungen für RAG mit langen Dokumenten

Zusammenfassung

Retrieval-augmented Generation (RAG)-Systeme müssen die Granularität der Abfrage mit der kontextuellen Kohärenz in Einklang bringen – eine Herausforderung, die bestehende Methoden durch LLM-gesteuertes Chunking, einstufige Kontexterweiterung oder hierarchische Zusammenfassung angehen. Diese Ansätze sind in unterschiedlichem Maße auf kostspielige LLM-Aufrufe während der Indizierung oder Abfrage angewiesen, beschränken die Kontextaggregation auf eine einzige Granularitätsebene oder führen durch Zusammenfassung zu Informationsverlust. Wir stellen SproutRAG vor, ein aufmerksamkeitsgesteuertes hierarchisches RAG-Framework, das diesen Zielkonflikt angeht, indem es satzweise Chunks in zunehmend größere, aber semantisch kohärente Einheiten organisiert und dabei gelernte Aufmerksamkeitsverknüpfungen zwischen Sätzen nutzt, um einen binären Chunking-Baum aufzubauen. Im Gegensatz zu früheren Ansätzen, die auf externe LLMs, feste Kontexterweiterung oder verlustbehaftete Zusammenfassung angewiesen sind, lernt SproutRAG, welche Aufmerksamkeitsköpfe und -schichten die semantische Dokumentenstruktur am besten erfassen, und ermöglicht so eine Multi-Granularitäts-Abfrage ohne zusätzliche LLM-Aufrufe oder komprimierte Zusammenfassungen. Zur Abfragezeit verwendet SproutRAG eine hierarchische Beam-Suche, um Kandidaten auf mehreren Granularitätsebenen abzurufen und so die Relevanz über mehrere Sätze hinweg über die flache Abfrage hinaus zu erfassen. Das Framework wird end-to-end mit einem gemeinsamen Ziel trainiert, das sowohl die Einbettungen als auch die Baumstruktur verbessert. Experimente mit vier Benchmarks aus den Bereichen Wissenschaft, Recht und offene Domänen zeigen, dass SproutRAG die Informationseffizienz (IE) im Durchschnitt um 6,1 % gegenüber der stärksten Baseline verbessert. Der Code ist unter https://github.com/AmirAbaskohi/SproutRAG verfügbar.

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.