KohakuRAG: A simple RAG framework with hierarchical document indexing

O artigo apresenta o KohakuRAG, um framework RAG hierárquico que preserva a estrutura documental através de uma representação em árvore de quatro níveis, utiliza um planejador de consultas baseado em LLM e inferência em conjunto para melhorar a precisão e a estabilidade, alcançando o primeiro lugar no Desafio WattBot 2025 com pontuação de 0,861.

O essencial

Imagine que você tem uma biblioteca gigante com 32 livros técnicos muito densos sobre o consumo de energia de Inteligência Artificial. Alguém chega e faz uma pergunta super específica: "Quanto de água foi usada para treinar o modelo X?".

O problema é que os livros estão bagunçados, a linguagem é difícil e, às vezes, a resposta está escondida em um gráfico ou em uma frase pequena no meio de um parágrafo gigante. Se você pedir para um robô (uma Inteligência Artificial) ler tudo e responder, ele pode inventar coisas (alucinar), esquecer de dizer de onde tirou a informação ou ficar confuso.

Os autores do KohakuRAG criaram uma solução inteligente para esse problema. Vamos explicar como funciona usando analogias do dia a dia:

1. A Biblioteca Organizada (Indexação Hierárquica)

O Problema: A maioria dos sistemas tenta ler os livros cortando-os em pedaços aleatórios de tamanho fixo, como se você pegasse uma tesoura e cortasse um livro de capa dura ao meio, misturando o índice com o capítulo 3. Isso perde a estrutura.

A Solução KohakuRAG: Eles tratam os documentos como uma árvore genealógica ou um mapa de endereços.

• O Livro inteiro é a raiz.
• Os Capítulos são os galhos grandes.
• Os Parágrafos são os galhos menores.
• As Frases são as folhas.

Ao invés de cortar aleatoriamente, o sistema entende que uma frase pertence a um parágrafo, que pertence a um capítulo. Quando o robô precisa de uma informação, ele não busca apenas uma "fatia" solta; ele busca a folha exata, mas também traz o galho e o ramo para cima, para que o robô entenda o contexto completo. É como pedir para um bibliotecário: "Traga-me a página 42", mas ele traz também o capítulo inteiro para garantir que você não perca o sentido da história.

2. O Detetive com Múltiplas Lentes (Recuperação Multi-Consulta)

O Problema: Às vezes, você pergunta "Qual o PUE?" (um termo técnico), mas o livro usa a frase "Eficiência no Uso de Energia". O robô, sendo literal, não acha nada. É como procurar "batata" na prateleira de "papas".

A Solução KohakuRAG: Eles usam um Detetive Inteligente (um modelo de IA) antes de procurar.

• O Detetive recebe sua pergunta e pensa: "Espera, essa pessoa pode estar procurando por 'PUE', 'Eficiência Energética' ou 'Razão de Consumo'".
• Ele cria várias versões da mesma pergunta (como se tivesse várias lentes diferentes) e as envia para a biblioteca.
• Depois, ele junta todas as respostas. Se vários "detetives" encontraram o mesmo documento, esse documento ganha mais pontos e vai para o topo da lista. É como se várias pessoas buscassem algo no mesmo lugar; se todos apontam para o mesmo baú, provavelmente o tesouro está lá.

3. O Painel de Especialistas (Inferência em Conjunto)

O Problema: Se você pedir a um único robô para responder, ele pode ter um "dia ruim", ficar confuso ou inventar uma resposta porque não encontrou o que queria.

A Solução KohakuRAG: Eles não confiam em um único robô. Eles organizam uma reunião de especialistas.

• O sistema pede a resposta para o mesmo robô várias vezes (digamos, 9 vezes), cada vez com uma pequena variação.
• Se o robô diz "Não sei" (abstenção) em uma das vezes, mas os outros 8 dizem "É 500 litros", o sistema ignora o "Não sei" e segue a maioria.
• Se todos dizem "Não sei", o sistema admite honestamente que não sabe, em vez de inventar uma mentira. Isso evita que o robô alucine.

4. O Mecanismo de "Tente de Novo"

Às vezes, o robô diz "Não sei" porque a informação estava um pouco escondida. O sistema tem um botão de retry (tente de novo). Se o robô desistir, o sistema aumenta a busca, traz mais contexto e pergunta de novo. É como se você estivesse procurando algo no fundo de uma gaveta bagunçada; se não acha na primeira olhada, você tira mais coisas de cima e tenta de novo.

O Resultado?

Na competição WattBot 2025, onde o desafio era responder perguntas técnicas com precisão de 0,1% e citar exatamente de onde veio a informação, o KohakuRAG venceu.

Eles foram a única equipe que manteve o 1º lugar tanto na fase pública quanto na fase privada (que é como um "segredo" que só é revelado no final). Isso prova que o sistema deles não apenas "decorou" as perguntas, mas aprendeu a buscar e entender a informação de verdade, sendo robusto e confiável.

Resumo da Ópera:
O KohakuRAG é como ter um bibliotecário superorganizado, um detetive que sabe falar várias línguas técnicas e um conselho de especialistas que vota juntos para garantir que a resposta seja correta, honesta e bem fundamentada.

Resumo técnico

Resumo Técnico: KohakuRAG

1. Problema e Contexto

Os sistemas de Geração Aumentada por Recuperação (RAG) enfrentam desafios significativos quando aplicados a tarefas que exigem citações precisas e tolerância numérica estrita. O artigo identifica três limitações fundamentais nos métodos RAG padrão:

• Fragmentação Plana (Flat Chunking): A divisão de documentos em segmentos de tamanho fixo destrói a estrutura hierárquica do documento, dificultando o rastreamento preciso da proveniência (citação) e a compreensão semântica contextual.
• Desconexão de Vocabulário: Formulações de consulta únicas frequentemente falham em recuperar passagens relevantes devido a diferenças de vocabulário entre a pergunta do usuário e o documento fonte (ex.: "PUE" vs. "Power Usage Effectiveness").
• Instabilidade de Resposta: A inferência única de Grandes Modelos de Linguagem (LLMs) é estocástica, gerando variações no conteúdo da resposta e na seleção de citações entre execuções, além de uma tendência a "abster-se" (não responder) desnecessariamente quando a evidência é difícil de localizar.

O trabalho foi motivado pelo WattBot 2025 Challenge, um benchmark que exige que os sistemas respondam a ~300 perguntas técnicas sobre consumo de energia de IA, com tolerância numérica de ±0,1%, citação exata de documentos e capacidade de abstenção correta quando a evidência é insuficiente.

2. Metodologia: KohakuRAG

O KohakuRAG propõe um framework RAG hierárquico que aborda os problemas acima através de três mecanismos principais:

A. Indexação Hierárquica de Documentos (Hierarchical Document Indexing)

• Estrutura em Árvore: Em vez de chunks planos, os documentos são parseados em uma árvore de quatro níveis: Documento → Seção → Parágrafo → Frase.
• Propagação de Embeddings Bottom-Up: Os embeddings são calculados nas folhas (frases) e agregados para os níveis superiores (parágrafos, seções, documento) usando uma média ponderada pelo número de tokens. Isso preserva a semântica composicional e permite limites de citação naturais em qualquer granularidade.
• Multimodalidade: Figuras e tabelas são tratadas como nós especiais. O sistema gera legendas técnicas usando modelos de visão (VLMs) e as integra na estrutura hierárquica, permitindo recuperação baseada em texto para informações visuais.

B. Recuperação com Múltiplas Consultas e Reordenamento (Multi-Query Retrieval)

• Planejador de Consultas (Query Planner): Um LLM gera múltiplas consultas semanticamente relacionadas a partir da pergunta original, cobrindo diferentes formulações, expandindo abreviações e decompondo perguntas compostas.
• Reordenamento Cruzado (Cross-Query Reranking): Os resultados de todas as consultas são agregados. Nós recuperados por múltiplas consultas recebem prioridade (sinal de consenso), superando o problema de vocabulário.
• Expansão de Contexto Hierárquico: Ao recuperar um nó (ex.: uma frase), o sistema automaticamente inclui seus nós pais (parágrafo, seção) para fornecer contexto suficiente ao LLM, evitando a necessidade de recuperar passagens redundantes.

C. Inferência de Ensemble com Filtragem de Abstenção (Ensemble Inference)

• Múltiplas Execuções: O sistema executa a geração de resposta $m$ vezes independentemente.
• Votação Majoritária com Tratamento de "Blank": Um mecanismo crucial é filtrar respostas em branco (abstenção) antes da votação, desde que existam respostas não vazias. Isso impede que execuções conservadoras dominem quando a evidência existe, mas é difícil de encontrar.
• Mecanismo de Retry: Se o modelo indicar abstenção (evidência insuficiente), o sistema aumenta o parâmetro $k$ (número de resultados recuperados) e tenta novamente, recuperando informações que estavam logo abaixo do corte inicial.

3. Contribuições Principais

1. Esquema de Indexação Hierárquica: Preserva a estrutura do documento e permite rastreamento de citação preciso através de agregação de embeddings bottom-up.
2. Planejamento de Consultas com LLM: Melhora a cobertura de recuperação expandindo a consulta e utilizando sinais de consenso para reordenar resultados.
3. Mecanismo de Ensemble com Consciência de Abstenção: Agrega múltiplas execuções filtrando abstenções desnecessárias, que foram identificadas como o modo de erro dominante (26,8% das falhas).
4. Análise Empírica Robusta: Demonstra que a recuperação densa hierárquica sozinha é altamente competitiva, e que otimizações de prompt e mecanismos de retry têm impacto maior do que a adição de busca híbrida (BM25).

4. Resultados e Desempenho

O sistema foi avaliado no WattBot 2025 Challenge:

• Classificação: KohakuRAG alcançou 1º lugar em ambas as divisões (pública e privada) do leaderboard, com uma pontuação final de 0,861.
• Estabilidade: Foi a única equipe a manter a posição de topo em ambas as partições, demonstrando robustez contra overfitting em dados públicos.
• Estudos de Ablação (Impacto Relativo):
  • Reordenação de Prompt (Contexto antes da Pergunta): +80% de melhoria relativa.
  • Mecanismo de Retry: +69% de melhoria relativa (especialmente em profundidades de recuperação baixas).
  • Votação de Ensemble com Filtragem de Blank: +1,2 pontos percentuais (pp).
  • Recuperação Densa vs. Híbrida: A recuperação densa hierárquica foi suficiente; a adição de BM25 (busca esparsa) trouxe apenas +3,1 pp de melhoria, sugerindo retornos decrescentes para métodos híbridos quando a estrutura de recuperação é rica.
• Análise de Erros: Os principais erros foram abstenção desnecessária (26,8%), incompatibilidade de citação (23,6%) e seleção de valor incorreta (22,2%). O sistema mitigou eficazmente a abstenção desnecessária.

5. Significância e Conclusão

O KohakuRAG demonstra que, para tarefas de QA baseadas em documentos técnicos com requisitos rigorosos de citação e precisão numérica:

• A estrutura do documento é tão importante quanto a semântica do texto; preservar a hierarquia natural (seções, parágrafos) supera a fragmentação plana.
• A robustez é alcançada através de ensembles e mecanismos de recuperação iterativa, não apenas através de modelos LLM maiores.
• A otimização do fluxo de inferência (ordem do prompt, tratamento de abstenção) pode ter um impacto mais significativo do que a complexidade do modelo de recuperação (ex.: adicionar BM25).

O código e os dados do framework foram liberados como software de código aberto, estabelecendo um novo baseline para sistemas RAG de alta precisão em benchmarks de citação.