SproutRAG : Recherche arborescente guidée par l’attention avec embeddings progressifs pour le RAG sur documents longs

Résumé

Les systèmes de génération augmentée par récupération (RAG) doivent concilier granularité de la récupération et cohérence contextuelle, un défi que les méthodes existantes abordent par le découpage guidé par LLM, l'expansion contextuelle à un seul niveau ou le résumé hiérarchique. Ces approches dépendent diversement d'appels coûteux au LLM lors de l'indexation ou de la récupération, limitent l'agrégation contextuelle à un seul niveau de granularité, ou introduisent une perte d'information via le résumé. Nous présentons SproutRAG, un cadre hiérarchique RAG guidé par l'attention qui traite ce compromis en organisant les blocs au niveau des phrases en unités progressivement plus grandes mais sémantiquement cohérentes, en utilisant l'attention inter-phrases apprise pour construire un arbre de découpage binaire. Contrairement aux approches antérieures qui reposent sur des LLM externes, une expansion contextuelle fixe ou un résumé avec perte, SproutRAG apprend quelles têtes d'attention et quelles couches capturent le mieux la structure documentaire sémantique, permettant une récupération multi-granularité sans appels supplémentaires au LLM ni résumés compressés. Lors de la récupération, SproutRAG utilise une recherche par faisceau hiérarchique pour récupérer des candidats à plusieurs granularités, capturant la pertinence multi-phrases au-delà de la récupération plate. Le cadre est entraîné de bout en bout avec un objectif conjoint qui améliore à la fois les plongements et la structure de l'arbre. Des expériences sur quatre benchmarks couvrant des domaines scientifiques, juridiques et ouverts montrent que SproutRAG améliore l'efficacité informationnelle (EI) de 6,1 % en moyenne par rapport à la baseline la plus forte. Le code est disponible sur 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.