KohakuRAG: A simple RAG framework with hierarchical document indexing

KohakuRAG es un marco de generación aumentada por recuperación (RAG) jerárquico que preserva la estructura de los documentos mediante un árbol de cuatro niveles, mejora la cobertura de recuperación con un planificador de consultas impulsado por LLM y estabiliza las respuestas mediante inferencia de conjunto, logrando el primer lugar en el Desafío WattBot 2025 al superar los desafíos de precisión numérica y atribución de fuentes.

Lo esencial

¡Hola! Imagina que tienes una biblioteca gigante con 32 libros técnicos muy densos sobre el consumo de energía de la Inteligencia Artificial. Ahora, alguien te hace una pregunta muy específica, como "¿Cuánta agua se gasta entrenando un modelo específico?". Tu trabajo es encontrar la respuesta exacta en esos libros, citar la página correcta y decirte si no sabes la respuesta en lugar de inventar una.

Este es el desafío que enfrentó el equipo Kohaku-Lab en una competencia llamada WattBot 2025. Ganaron el primer lugar con su sistema llamado KohakuRAG.

Aquí te explico cómo funciona su sistema usando analogías sencillas:

1. El Problema: La Biblioteca Desordenada

La mayoría de los sistemas de inteligencia artificial (IA) que leen documentos hacen algo llamado "trocear" (chunking). Imagina que toman un libro, lo cortan en trozos de papel de 500 palabras al azar y los tiran en una pila.

• El problema: Si cortas un párrafo por la mitad, pierdes el contexto. Además, si buscas "eficiencia energética" pero el libro dice "uso de energía", el robot no encuentra nada porque son palabras diferentes. Y si el robot adivina mal, inventa datos (alucinaciones).

2. La Solución de KohakuRAG: Un Organizador de Mentes Maestras

El equipo creó un sistema con tres trucos principales para ganar:

A. El Árbol de Conocimiento (Indexación Jerárquica)

En lugar de cortar los libros en trozos aleatorios, KohakuRAG construye un árbol genealógico de la información.

• La analogía: Imagina que el documento es un árbol gigante.
  • La raíz es el libro entero.
  • Las ramas grandes son los capítulos.
  • Las ramas pequeñas son los párrafos.
  • Las hojas son las oraciones individuales.
• Cómo ayuda: El sistema lee desde las hojas hacia arriba. Si encuentras una hoja (una oración) importante, el sistema sabe automáticamente a qué rama (párrafo) y a qué árbol (capítulo) pertenece. Esto le permite decirte exactamente: "La respuesta está en la página 42, párrafo 3", sin perderse.

B. El Detective con Lupa Múltiple (Búsqueda Multi-Pregunta)

A veces, tú preguntas de una forma y el libro habla de otra.

• La analogía: Imagina que un detective (una IA muy inteligente) recibe tu pregunta. En lugar de ir a buscar con una sola frase, el detective imagina 4 o 5 versiones diferentes de tu pregunta.
  • Si preguntas: "¿Qué es el PUE?", el detective también busca: "¿Qué es la eficiencia en el uso de energía?", "¿Cómo se mide el consumo en centros de datos?", etc.
• El filtro de consenso: Luego, el detective reúne todas las respuestas de sus diferentes búsquedas. Si un párrafo aparece en varias de las búsquedas diferentes, ¡ese párrafo es muy importante! El sistema le da prioridad a esa información.

C. El Jurado de 10 Personas (Inferencia por Conjunto)

Las IAs a veces son nerviosas; si les preguntas lo mismo dos veces, pueden dar respuestas ligeramente distintas.

• La analogía: En lugar de confiar en la opinión de una sola persona, KohakuRAG convoca a un jurado de 10 personas (ejecuta el sistema 10 veces).
  • Si 8 personas dicen "La respuesta es 500 litros" y 2 dicen "No sé", el sistema ignora a las 2 que dicen "No sé" (porque quizás solo no encontraron la información esa vez) y sigue a la mayoría.
  • Si todas dicen "No sé", entonces el sistema es honesto y dice: "No tengo la información", en lugar de inventar una respuesta falsa.

3. El Truco Secreto: El Orden de las Cosas

El equipo descubrió algo curioso: a las IAs les cuesta más recordar lo que está en el medio de un texto largo (como si estuvieras leyendo un libro y olvidas lo que pasó en el capítulo 5 porque estás en el 10).

• La solución: Ponen la pregunta y las instrucciones al final, después de haberles dado todos los documentos. Así, la IA tiene la pregunta fresca en su mente justo cuando necesita responder.

4. El Resultado: Ganar con Consistencia

En la competencia, había dos listas de preguntas: una pública (que todos veían) y una privada (secreta hasta el final).

• Muchos equipos hicieron trampa (sin querer) estudiando solo la lista pública y luego fallaron en la privada.
• KohakuRAG fue el único equipo que mantuvo el primer lugar en ambas listas.
• ¿Por qué? Porque su sistema no solo era "listo", sino robusto. No dependía de adivinar, sino de buscar bien, verificar varias veces y ser honesto cuando no sabía.

En resumen

KohakuRAG es como un bibliotecario experto que:

1. No corta los libros, sino que entiende su estructura (capítulos y párrafos).
2. Pide a un detective que busque la información usando muchas palabras diferentes.
3. Consulta a un jurado entero para asegurarse de que la respuesta es correcta y no una alucinación.
4. Es tan honesto que prefiere decir "no sé" antes que inventar un dato.

¡Y así ganaron la competencia!

Resumen técnico

Resumen Técnico: KohakuRAG

1. El Problema

Los sistemas de Generación Aumentada por Recuperación (RAG) enfrentan dificultades críticas cuando se requieren citas de alta precisión y respuestas numéricas exactas. Los enfoques estándar presentan tres limitaciones fundamentales:

• Pérdida de estructura: Las estrategias de "chunking" (fragmentación) plana destruyen la jerarquía natural de los documentos (secciones, párrafos), complicando el rastreo de la procedencia exacta de la información.
• Desajuste de vocabulario: Las consultas únicas a menudo fallan al recuperar pasajes relevantes si el usuario utiliza terminología diferente a la del documento (ej. preguntar por "PUE" cuando el texto habla de "eficacia en el uso de la energía").
• Inestabilidad en la inferencia: Una sola ejecución de un Modelo de Lenguaje Grande (LLM) produce respuestas estocásticas que varían en contenido y selección de citas entre ejecuciones, y a menudo se abstienen innecesariamente cuando la evidencia es difícil de localizar.

Estos problemas se hacen evidentes en desafíos como el WattBot 2025, que exige responder preguntas técnicas sobre consumo energético de IA con una tolerancia numérica de $\pm0.1\%$, citas exactas de fuentes y la capacidad de abstenerse correctamente cuando la evidencia es insuficiente.

2. Metodología: KohakuRAG

KohakuRAG es un marco de trabajo RAG jerárquico diseñado para abordar estos desafíos mediante tres mecanismos principales:

A. Indexación Jerárquica de Documentos
En lugar de fragmentos planos, el sistema representa los documentos como árboles de cuatro niveles:

1. Documento $\rightarrow$ Sección $\rightarrow$ Párrafo $\rightarrow$ Oración.

• Propagación de Embeddings Bottom-Up: Se calculan los embeddings de las oraciones (hojas) y se agregan hacia arriba a los niveles de párrafo y sección mediante un promedio ponderado por longitud. Esto captura la semántica composicional y permite citas precisas en cualquier nivel de granularidad.
• Multimodalidad: Las imágenes y tablas se tratan como nodos especiales. Se utilizan modelos de visión (Qwen-VL) para generar descripciones (captions) que se integran en la estructura del árbol, permitiendo la recuperación basada en texto de información visual.

B. Recuperación Multi-Consulta con Reordenamiento (Reranking)
Para superar la brecha de vocabulario:

• Planificador de Consultas (Query Planner): Un LLM genera múltiples consultas semánticamente relacionadas a partir de la pregunta original (reformular, expandir acrónimos, descomponer preguntas compuestas).
• Recuperación y Reordenamiento: Se recuperan los resultados para cada consulta y se aplican estrategias de reordenamiento basadas en el consenso:
  • Frecuencia: Los nodos recuperados por más consultas reciben mayor rango.
  • Puntuación Combinada: Se normaliza y combina la frecuencia con la puntuación de similitud total.
• Expansión de Contexto Jerárquico: Al recuperar un nodo (ej. una oración), el sistema incluye automáticamente sus nodos padres (párrafo, sección) para proporcionar contexto suficiente al LLM sin redundancia.

C. Inferencia de Conjunto (Ensemble) con Votación Consciente de la Abstención
Para mitigar la inestabilidad:

• Se ejecutan $m$ inferencias independientes.
• Filtrado de Respuestas en Blanco: Si una ejecución indica "insuficiente evidencia" (abstención), pero existen otras que dan una respuesta, las respuestas en blanco se filtran antes de la votación. Esto evita que ejecuciones conservadoras dominen cuando la evidencia existe pero es difícil de encontrar.
• Votación Mayoritaria: Se selecciona la respuesta y las citas más frecuentes entre las ejecuciones válidas.

3. Contribuciones Clave

1. Esquema de Indexación Jerárquica: Preserva las relaciones estructurales mediante representaciones de árbol con agregación de embeddings, permitiendo un seguimiento preciso de la procedencia.
2. Planificador de Consultas Potenciado por LLM: Mejora la cobertura de recuperación expandiendo la pregunta y utilizando señales de consenso para el reordenamiento.
3. Mecanismo de Votación con Filtrado de Abstención: Aborda el modo de error dominante (abstención innecesaria) agregando múltiples ejecuciones y filtrando las respuestas vacías cuando hay consenso de respuesta.
4. Hallazgos Empíricos: Demuestra que la recuperación densa jerárquica por sí sola es altamente competitiva, y que la ordenación del prompt y los mecanismos de reintento aportan más valor que las estrategias híbridas complejas.

4. Resultados y Evaluación

El sistema fue evaluado en el WattBot 2025 Challenge (32 documentos, ~500k tokens, ~300 preguntas técnicas).

• Rendimiento: KohakuRAG obtuvo el primer lugar en ambos tableros de clasificación (público y privado) con una puntuación final de 0.861. Fue el único equipo que mantuvo la posición #1 en ambas particiones.
• Estudios de Ablación:
  • Orden del Prompt: Colocar el contexto antes de la pregunta (C$\rightarrow$Q) mejoró el rendimiento relativo en un +80%, confirmando el fenómeno "perdido en el medio" (lost in the middle).
  • Mecanismo de Reintento: Aumentar el contexto ($k$) cuando el modelo se abstiene proporcionó una mejora del +69% en profundidad de recuperación baja.
  • Votación de Conjunto: El filtrado de respuestas en blanco añadió +1.2 puntos porcentuales a la puntuación final.
  • Recuperación Híbrida: La adición de BM25 (búsqueda esparcida) solo aportó un +3.1pp, sugiriendo que una estructura de recuperación densa rica es suficiente para tareas con muchas citas.
• Análisis de Errores: El error más común (26.8%) fue la abstención innecesaria, seguida de la mismatch de citas (23.6%) y errores de selección de valor (22.2%).

5. Significado e Impacto

KohakuRAG demuestra que para tareas de preguntas y respuestas basadas en documentos técnicos que requieren alta precisión y citas:

1. La estructura importa más que la complejidad: Mantener la jerarquía del documento es más efectivo que simplemente aumentar la complejidad de la recuperación híbrida.
2. La robustez supera al pico de rendimiento: Los enfoques de conjunto (ensembles) y la gestión inteligente de la abstención son cruciales para generalizar bien a datos no vistos (partición privada), evitando el sobreajuste a las características de un subconjunto de datos específico.
3. Eficiencia en la inferencia: Pequeños ajustes en el diseño del prompt y mecanismos de reintento tienen un impacto mayor que el uso de modelos más grandes o estrategias de recuperación excesivamente complejas.

El código y el marco de trabajo se han liberado como software de código abierto, proporcionando una base sólida para sistemas RAG que requieren trazabilidad y precisión en entornos técnicos.