SproutRAG: Busca em Árvore Guiada por Atenção com Embeddings Progressivos para RAG de Documentos Longos

Resumo

Os sistemas de geração aumentada por recuperação (RAG) precisam equilibrar a granularidade da recuperação com a coerência contextual, um desafio que os métodos existentes abordam por meio de fragmentação guiada por LLM, expansão de contexto em nível único ou sumarização hierárquica. Essas abordagens dependem, de diferentes formas, de chamadas custosas ao LLM durante a indexação ou recuperação, limitam a agregação de contexto a um único nível de granularidade ou introduzem perda de informações por meio da sumarização. Apresentamos o SproutRAG, uma estrutura hierárquica de RAG guiada por atenção que aborda essa troca ao organizar fragmentos em nível de sentença em unidades progressivamente maiores, mas semanticamente coerentes, usando a atenção aprendida entre sentenças para construir uma árvore binária de fragmentação. Diferentemente de abordagens anteriores que dependem de LLMs externos, expansão fixa de contexto ou sumarização com perda, o SproutRAG aprende quais cabeças e camadas de atenção melhor capturam a estrutura semântica do documento, permitindo recuperação em múltiplas granularidades sem chamadas adicionais ao LLM ou sumários comprimidos. No momento da recuperação, o SproutRAG usa busca em feixe hierárquica para recuperar candidatos em múltiplas granularidades, capturando relevância multissentença além da recuperação plana. A estrutura é treinada de ponta a ponta com um objetivo conjunto que melhora tanto os embeddings quanto a estrutura da árvore. Experimentos em três conjuntos de referência abrangendo domínios científicos, jurídicos e de código aberto demonstram que o SproutRAG melhora a eficiência da informação (EI) em 6,1% em média em relação ao modelo de base mais forte. O código está disponível em 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.