SproutRAG: Búsqueda en Árbol Guiada por Atención con Embeddings Progresivos para RAG en Documentos Largos

Resumen

Los sistemas de generación aumentada por recuperación (RAG) deben equilibrar la granularidad de la recuperación con la coherencia contextual, un desafío que los métodos existentes abordan mediante la fragmentación guiada por LLM, la expansión de contexto de un solo nivel o el resumen jerárquico. Estos enfoques dependen de costosas llamadas a LLM durante la indexación o la recuperación, limitan la agregación de contexto a un solo nivel de granularidad o introducen pérdida de información mediante el resumen. Presentamos SproutRAG, un marco RAG jerárquico guiado por atención que aborda esta compensación organizando fragmentos a nivel de oración en unidades progresivamente más grandes pero semánticamente coherentes, utilizando la atención entre oraciones aprendida para construir un árbol de fragmentación binario. A diferencia de enfoques anteriores que dependen de LLM externos, expansión de contexto fija o resumen con pérdida, SproutRAG aprende qué cabezas y capas de atención capturan mejor la estructura semántica del documento, permitiendo la recuperación de múltiples granularidades sin llamadas adicionales a LLM ni resúmenes comprimidos. En el momento de la recuperación, SproutRAG utiliza búsqueda de haz jerárquica para recuperar candidatos en múltiples granularidades, capturando relevancia entre oraciones más allá de la recuperación plana. El marco se entrena de extremo a extremo con un objetivo conjunto que mejora tanto las incrustaciones como la estructura del árbol. Experimentos en cuatro puntos de referencia que abarcan entornos científicos, legales y de dominio abierto demuestran que SproutRAG mejora la eficiencia de información (EI) en un 6.1% en promedio sobre la línea base más fuerte. El código está disponible en 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.