SproutRAG: Поиск по дереву, управляемый вниманием, с прогрессивными вложениями для RAG длинных документов

Аннотация

Системы генерации с дополнением поиском (RAG) должны балансировать гранулярность поиска и контекстную связность — задача, которую существующие методы решают с помощью разбиения на фрагменты с использованием LLM, расширения контекста на одном уровне или иерархического обобщения. Эти подходы в разной степени зависят от дорогостоящих вызовов LLM на этапах индексации или поиска, ограничивают агрегацию контекста одним уровнем гранулярности или приводят к потере информации при обобщении. Мы представляем SproutRAG — иерархический RAG-фреймворк, управляемый механизмом внимания, который устраняет этот компромисс, организуя фрагменты на уровне предложений в последовательно укрупняющиеся, но семантически связные единицы с использованием обученного межпредложенческого внимания для построения бинарного дерева фрагментации. В отличие от предыдущих подходов, опирающихся на внешние LLM, фиксированное расширение контекста или сжатое обобщение, SproutRAG обучается определять, какие головы и слои внимания наилучшим образом отражают семантическую структуру документа, что обеспечивает многогранулярный поиск без дополнительных вызовов LLM или сжатых изложений. На этапе поиска SproutRAG использует иерархический лучевой поиск для извлечения кандидатов на нескольких уровнях гранулярности, захватывая релевантность на уровне нескольких предложений за пределами плоского поиска. Фреймворк обучается сквозным образом с совместной целевой функцией, улучшающей как эмбеддинги, так и структуру дерева. Эксперименты на четырёх наборах данных, охватывающих научные, юридические и открытые домены, показывают, что SproutRAG в среднем улучшает информационную эффективность (ИЭ) на 6,1% по сравнению с наилучшим базовым методом. Код доступен по адресу 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.