Hierarchical Abstract Tree for Cross-Document Retrieval-Augmented Generation

1. Formulação do Problema

A Geração Aumentada por Recuperação (RAG) aprimora os Modelos de Linguagem de Grande Escala (LLMs) com conhecimento externo. Embora os métodos Tree-RAG (como o RAPTOR) organizem documentos em índices hierárquicos para suportar consultas de multi-granularidade, eles enfrentam limitações críticas ao escalar para questões multi-hop entre documentos:

1. Baixa Adaptabilidade à Distribuição: Os métodos existentes dependem de clustering do tipo $k$-means (por exemplo, Modelos de Mistura Gaussiana). Esses algoritmos assumem distribuições de dados esféricas e sofrem com o "efeito uniforme", forçando os clusters a terem tamanhos uniformes. Isso introduz ruído ao fundir tópicos menores e enviesados em clusters maiores, fazendo com que o recuperador ignore informações de nicho.
2. Isolamento Estrutural: Os índices em árvore geralmente carecem de conexões explícitas entre nós folha de diferentes ramos. Diferentemente do Graph-RAG, que salta entre entidades via arestas, as folhas da árvore estão isoladas. Isso impede que o recuperador capture dependências causais implícitas necessárias para o raciocínio multi-hop.
3. Abstração Grossa: A abstração padrão em árvores frequentemente cria resumos de alto nível que obscurecem detalhes finos. A correspondência de vetores densos luta para associar entidades específicas em uma consulta do usuário a esses conceitos abstratos, levando a falhas de recuperação em questões factuais ao nível de token.

2. Metodologia: Ψ-RAG

Os autores propõem o Ψ-RAG, um framework composto por dois componentes principais: um Índice de Árvore Abstrata Hierárquica e um Recuperador Agente Multi-granular.

A. Indexação de Árvore Abstrata Hierárquica

Em vez de clustering $k$-means, o Ψ-RAG utiliza um processo iterativo de "fusão e colapso" inspirado no Clustering Hierárquico Aglomerativo (AHC), mas sem pressupostos de distribuição a priori.

• Classificação de Similaridade: Todos os fragmentos de documento são codificados em vetores densos, e as similaridades pares são calculadas e classificadas.
• Fusão e Colapso:
  • Fusão: O par mais similar de fragmentos (ou nós) é vinculado sob um novo nó abstrato.
  • Colapso de Nó Folha: Se um nó já estiver vinculado a um pai, o nó isolado é anexado ao mesmo pai.
  • Colapso de Nó Abstrato: Se dois nós pertencem a pais diferentes na mesma profundidade, suas raízes são fundidas sob um novo nó. Se as profundidades diferirem, a raiz mais rasa é vinculada ao caminho do nó mais profundo para equilibrar a árvore.
• Abstração: Um agente LLM gera dois tipos de abstrações para cada nó não folha:
  • Abstração Sumativa: Um parágrafo coerente que resume a lógica e as relações entre entidades.
  • Abstração de Palavras-chave: Palavras-chave/frases de alta densidade que preservam a co-ocorrência de entidades.
• Reequilíbrio: Nós com filhos excessivos são divididos para garantir que a janela de contexto do agente de abstração não seja excedida.

B. Recuperação Agente Multi-granular

Para abordar o isolamento estrutural e a abstração grossa, o Ψ-RAG emprega um Agente de Recuperação e Resposta (R&A) que interage com um Índice Híbrido Multi-granular (Árvore + Índice de Palavras-chave Esparsas).

• Raciocínio Iterativo: O agente avalia se as informações atuais são suficientes. Se não forem, ele gera uma consulta reorganizada (enriquecida com contexto, por exemplo, adicionando aposições descritivas) e dispara uma nova etapa de recuperação.
• Recuperação Híbrida:
  • Recuperação em Árvore: Uma busca de vetor denso de cima para baixo, começando da raiz até a camada de folhas.
  • Recuperação Esparsa: Uma busca por palavras-chave baseada em BM25.
  • Fusão: Os resultados são combinados por meio de um reclassificador (paramétrico) ou Fusão de Classificação Recíproca (não paramétrico).
• Reorganização de Consulta: O agente reescreve dinamicamente as consultas para incluir palavras-chave temáticas, auxiliando tanto o recuperador esparso (encontrando entidades específicas) quanto o recuperador em árvore (localizando abstrações de alto nível relevantes).

3. Principais Contribuições

1. Adaptabilidade à Distribuição: O Ψ-RAG prova teoricamente (via análise de custo de Dasgupta) que seu processo de fusão hierárquica evita o "efeito uniforme" do $k$-means. Ele preserva adaptativamente distribuições de dados enviesadas, garantindo que clusters menores permaneçam distintos em vez de serem absorvidos pelos maiores.
2. Primeiro Tree-RAG entre Documentos: É o primeiro framework Tree-RAG capaz de lidar com recuperação multi-hop entre documentos em nível de corpus, igualando as capacidades do Graph-RAG enquanto mantém a eficiência das estruturas de árvore.
3. Framework Agente Híbrido: Resolve o conflito entre abstração grossa e recuperação de granularidade fina, integrando um índice de palavras-chave esparsas e um agente iterativo que reorganiza consultas com base nas necessidades de raciocínio.
4. Flexibilidade: O framework é construído inteiramente sobre LLMs de código aberto e é modular, permitindo que componentes (embedding, abstração, agente, reclassificador) sejam trocados sem retreinamento.

4. Resultados Experimentais

Os autores avaliaram o Ψ-RAG em diversos benchmarks, incluindo QA de salto único/multi-hop (HotpotQA, 2Wiki, MuSiQue, MultiHop-RAG), QA narrativa (NarrativeQA) e sumarização (QMSum, WCEP).

• Ganhos de Desempenho:
  • Em benchmarks de QA multi-hop entre documentos, o Ψ-RAG supera o RAPTOR em 25,9% no escore F1 médio.
  • Supera o HippoRAG 2 (um método Graph-RAG) mais avançado em 7,4% no escore F1 médio.
  • Alcança ganhos significativos em métricas de recuperação (Recall@5) em todos os conjuntos de dados.
• Eficiência:
  • Indexação: O Ψ-RAG constrói uma árvore em nível de corpus 10x mais rápida que o Graph-RAG baseado em OpenIE (por exemplo, GraphRAG) e 6,5x mais rápida que o RAPTOR.
  • Recuperação: A busca de cima para baixo oferece complexidade $O(\log n)$, resultando em menor latência que a travessia do Graph-RAG.
• Estudos de Ablação:
  • A remoção do agente R&A causa uma queda de ~20% no desempenho multi-hop.
  • A remoção do recuperador esparso prejudica significativamente o desempenho em consultas factuais (por exemplo, -21,8% no PopQA).
  • A reorganização de consultas pelo agente é crítica para o sucesso multi-hop.

5. Significado

• Ponte entre Lacunas: O Ψ-RAG conecta com sucesso a lacuna entre a eficiência do Tree-RAG e o poder de raciocínio do Graph-RAG, provando que estruturas hierárquicas podem lidar efetivamente com raciocínio multi-hop complexo sem a sobrecarga pesada da construção de grafos.
• Escalabilidade: Ao evitar o "efeito uniforme" e utilizar algoritmos de fusão eficientes, ele escala efetivamente para grandes corpora enviesados onde o clustering tradicional falha.
• Praticidade: O framework demonstra que modelos de código aberto podem alcançar desempenho de última geração em RAG estruturado, tornando o raciocínio multi-hop avançado acessível sem APIs proprietárias ou retreinamento extensivo.
• Direção Futura: Estabelece um novo padrão para o tratamento de recuperação entre documentos, sugerindo que a indexação hierárquica combinada com raciocínio agente é um caminho viável e superior para sistemas RAG de próxima geração.