Resumo Técnico: SPROUTRAG

Definição do Problema

Sistemas de geração aumentada por recuperação (RAG) enfrentam um compromisso fundamental entre granularidade de recuperação e coerência contextual. Os métodos existentes lutam para equilibrar essas necessidades conflitantes:

• Chunks grandes preservam o contexto, mas introduzem ruído redundante que dilui informações cruciais.
• Chunks de granularidade fina oferecem precisão, mas sofrem de fragmentação semântica e quebra de relações entre os chunks, particularmente em tarefas de recuperação cross-paragraph (ex: raciocínio multi-hop).
• Soluções atuais frequentemente dependem de chamadas custosas de Large Language Models (LLMs) durante a indexação ou recuperação (ex: chunking guiado por LLM, expansão de pares com filtragem por LLM), limitam a agregação de contexto a um único nível de granularidade ou introduzem perda de informação através de sumarização hierárquica (ex: RAPTOR).

O artigo argumenta que há necessidade de um framework que possa organizar a estrutura do documento para suportar recuperação de multi-granularidade sem incorrer no overhead de inferência de LLMs externos ou na perda de informação da sumarização.

Metodologia: SPROUTRAG

SPROUTRAG (Attention-Guided Tree Search with Progressive Embeddings) é um framework de RAG hierárquico que constrói uma árvore binária aprendida sobre chunks ao nível de sentença. Ele opera em duas fases:

1. Indexação Guiada por Atenção (Offline)

• Codificação ao Nível de Sentença: Um Modelo de Linguagem de Nível de Sentença (SLLM) codifica o documento para produzir embeddings de sentenças contextualizados e matrizes de atenção inter-sentença através de todas as camadas e cabeças do transformer.
• Agregação de Atenção Aprendida: Para abordar o "viés de proximidade" inerente à média uniforme de cabeças de atenção (que enfatiza excessivamente a adjacência sequencial), o SPROUTRAG aprende uma agregação ponderada de cabeças e camadas de atenção. Pesos escalares aprendíveis ($\alpha_{l,h}$) permitem que o modelo identifique quais cabeças de atenção específicas melhor capturam a correlação semântica para a estrutura do documento.
• Construção da Árvore: Uma árvore binária é construída de baixo para cima. Inicialmente, cada sentença é uma folha. Em cada etapa, o par de nós ativos com a maior pontuação de atenção mútua é mesclado.
• Embeddings Progressivos: Nós internos armazenam "embeddings progressivos" computados como a média dos embeddings de seus filhos ($e_p = (e_u + e_v)/2$). Isso permite que os nós internos representem unitariamente unidades semânticas mais amplas.
• Atualização de Single-Linkage: Ao mesclar nós, o pai herda a relação mais forte com os nós restantes, preservando conexões semânticas de longo alcance conforme a hierarquia cresce.

2. Recuperação Hierárquica (Online)

• Codificação da Query: A query é codificada usando o mesmo SLLM.
• Busca em Árvore (Beam Search) Hierárquica: A recuperação atravessa a árvore usando uma estratégia de beam search. Começando da raiz, o algoritmo expande os filhos dos nós atuais do beam e retém os top-$b$ nós mais relevantes baseados em similaridade de cosseno.
• Coleta de Multi-Granularidade: Diferente da recuperação plana, o SPROUTRAG coleta candidatos de todos os níveis de travessia da árvore (folhas, nós internos e subárvores) que excedem um limiar de similaridade aprendido.
• Reranking: Os candidatos coletados são reclassificados, e os top-$k$ chunks são passados para o gerador. Isso permite que o sistema recupere fatos precisos (folhas) ou contextos mais amplos de múltiplas sentenças (nós internos), conforme apropriado para a query.

Objetivo de Treinamento Conjunto

O framework é treinado de ponta a ponta com uma função de perda conjunta:

1. Objetivo de Recuperação ($L_{ret}$): Uma perda de aprendizado contrastivo que alinha os embeddings da query com passagens positivas e os separa de negativos difíceis (hard negatives).
2. Objetivo de Estrutura ($L_{attn}$): Um regularizador de atenção que encoraja pares de sentenças co-relevantes (aqueles que apoiam conjuntamente uma query) a receberem alta atenção mútua, moldando diretamente a estrutura da árvore induzida.

Principais Contribuições

1. Framework SPROUTRAG: Um sistema RAG hierárquico guiado por atenção que constrói uma árvore binária sobre chunks ao nível de sentença usando atenção inter-sentença aprendida, permitindo recuperação de multi-granularidade sem qualquer chamada de LLM no tempo de inferência.
2. Endereçamento do Viés de Proximidade: A identificação e mitigação do viés de proximidade em transformers de nível de sentença. O artigo substitui a média uniforme de atenção ingênua por uma agregação ponderada aprendida, permitindo que o modelo descubra quais cabeças de atenção melhor refletem a correlação semântica para a estrutura do documento.
3. Treinamento Conjunto: Um objetivo de treinamento conjunto inovador que otimiza simultaneamente a qualidade da recuperação e a estrutura da árvore, eliminando a necessidade de filtragem externa por LLM ou sumarização com perda de informação em qualquer estágio do pipeline.

Resultados Experimentais

Os autores avaliaram o SPROUTRAG em quatro benchmarks abrangendo configurações científicas (SCI-DOCS), jurídicas (LegalBench-RAG), de domínio aberto (MS MARCO) e multi-domínio (Dragonball), bem como tarefas de geração de ponta a ponta (HotpotQA, WebQuestions).

• Desempenho de Recuperação: O SPROUTRAG alcançou a maior Eficiência de Informação (IE) (Recall $\times$ Precisão) em todos os datasets, melhorando em média 6,1% sobre o baseline mais forte. Especificamente, melhorou a IE em 8,06 pontos no Dragonball, 4,65 no SCI-DOCS, 4,90 no LegalBench-RAG e 6,83 no MS MARCO.
• Precisão vs. Recall: Ao contrário de métodos que sacrificam a precisão pela recall (ou vice-versa), o SPROUTRAG alcançou a melhor precisão em cada benchmark enquanto mantinha um alto recall. Isso sugere que a hierarquia recupera efetivamente o contexto de suporte mais amplo sem introduzir ruído.
• Eficiência: Em tarefas de geração de ponta a ponta, o SPROUTRAG ofereceu o trade-off desempenho-eficiência mais forte. Superou sistemas pesados baseados em raciocínio (como GraphRAG e PageIndex) enquanto utilizava significativamente menos tokens online por query (4.38K) e menor latência (193 ms).
• Estudos de Ablação:
  • Remover os objetivos de treinamento conjunto degradou significativamente o desempenho, confirmando a necessidade tanto do alinhamento de recuperação quanto da regularização de estrutura.
  • Substituir a agregação de atenção aprendida por média uniforme ou árvores de similaridade de embedding resultou em menor desempenho, validando a importância dos sinais de atenção aprendidos para a construção da árvore.
  • A recuperação apenas de folhas manteve alta precisão, mas sofreu em recall e IE, demonstrando o valor de recuperar nós internos.

Significância e Alegações

O artigo afirma que o SPROUTRAG aborda com sucesso o compromisso entre granularidade de recuperação e coerência contextual. Ao organizar chunks ao nível de sentença em uma hierarquia aprendida e guiada por atenção, o framework permite a recuperação de multi-granularidade que captura a relevância emergente de múltiplas sentenças que métodos planos ou de pares perdem.

Crucialmente, os autores enfatizam que isso é alcançado sem chamadas de LLM no tempo de inferência ou sumarização com perda de informação. O framework generaliza através de diversos domínios (científico, jurídico, aberto) após um estágio de treinamento único, amortizando o custo de treinamento. Os resultados sugerem que sinais de atenção aprendidos podem efetivamente modelar a estrutura do documento para recuperação, oferecendo um caminho para sistemas RAG eficientes e de alta qualidade que rivalizam com abordagens mais computacionalmente caras e dependentes de LLM.

O artigo reconhece limitações, observando que a atual estrutura de árvore binária pode ser restritiva para dependências muitos-para-muitos e que o custo de treinamento inicial é maior do que usar recuperadores prontos para uso sem adaptação. No entanto, posiciona o SPROUTRAG como um passo significativo em direção à recuperação eficiente e consciente da estrutura.