KohakuRAG: A simple RAG framework with hierarchical document indexing

KohakuRAG is een hiërarchisch RAG-framework dat documentstructuur behoudt via een boomrepresentatie, zoekopdrachten verbetert met een LLM-geplande planner en antwoorden stabiliseert door ensemble-inferentie, waarmee het de eerste plaats behaalt op de WattBot 2025 Challenge-benchmark.

De kern

Stel je voor dat je een enorme bibliotheek binnenstapt, vol met 32 zeer technische boeken over de energieverbruik van kunstmatige intelligentie. Je krijgt een vraag: "Hoeveel water verbruikt een AI-model tijdens het trainen?" en je moet het antwoord geven met een precisie van 0,1%, terwijl je ook exact moet aangeven op welke bladzijde je het hebt gevonden. Als je het antwoord niet weet, mag je eerlijk zeggen: "Ik weet het niet."

Dit is precies de uitdaging waar KohakuRAG voor is ontworpen. Het is een slimme "zoekmachine en vertaler" die helpt om vragen te beantwoorden op basis van grote documenten, zonder dat de computer verzint (hallucineert) of de bron verliest.

Hier is hoe KohakuRAG werkt, vertaald naar alledaagse taal:

1. De Bibliotheek niet in stukjes hakken, maar in lagen ordenen

Het oude probleem:
Stel je voor dat je een boek in de schuur gooit en het in duizenden willekeurige stukjes papier snijdt. Als je nu een vraag stelt, zoekt de computer in die losse flarden. Het probleem? De context is weg. Je weet niet meer welke paragraaf bij welk hoofdstuk hoorde. En als je een antwoord vindt, weet je niet meer of het uit hoofdstuk 1 of 10 komt.

De KohakuRAG-oplossing:
KohakuRAG behandelt documenten als een groot, goed georganiseerd huis.

• Het dak is het hele document.
• De kamers zijn de hoofdstukken.
• De muren zijn de alinea's.
• De tegels zijn de zinnen.

In plaats van alles in één grote hoop te gooien, bouwt het systeem een hiërarchische kaart. Als het een zin vindt, weet het direct: "Ah, deze zin zit in de badkamer van de eerste verdieping." Dit maakt het veel makkelijker om later te zeggen: "Het antwoord staat op pagina 12, in de paragraaf over waterverbruik."

2. De slimme vertaler die meerdere vragen stelt

Het oude probleem:
Soms zoekt de computer op het woord "stroomverbruik", maar staat het antwoord in het boek onder "energie-efficiëntie". Een simpele zoekopdracht mist dan het antwoord, omdat de woorden niet exact overeenkomen.

De KohakuRAG-oplossing:
KohakuRAG gebruikt een slimme vertaler (een AI) die als een detective werkt. Als jij vraagt: "Hoeveel stroom verbruikt Google?", denkt de vertaler niet alleen aan "stroom", maar bedenkt hij ook:

• "Misschien staat het onder 'energie-efficiëntie'?"
• "Of misschien onder 'PUE' (een technische afkorting)?"
• "Of misschien in een rapport over 'duurzaamheid'?"

De vertaler stelt dus vier verschillende vragen tegelijk aan de bibliotheek. Vervolgens kijkt hij welke antwoorden door meerdere vragen werden gevonden. Als drie van de vier zoektochten naar dezelfde pagina wijzen, weet hij: "Dit is waarschijnlijk het juiste antwoord!"

3. Het panel van experts (Ensemble)

Het oude probleem:
Als je één persoon vraagt om een moeilijk vraag te beantwoorden, kan die persoon een slechte dag hebben, een fout maken of twijfelen. In de wereld van AI noemen we dit "stochastisch" (willekeurig). Soms geeft de computer een goed antwoord, soms een raar antwoord, en soms zegt hij "Ik weet het niet", terwijl het antwoord er wel was.

De KohakuRAG-oplossing:
KohakuRAG roept negen verschillende experts (of dezelfde computer, maar negen keer) om het antwoord te geven.

• Als 8 experts zeggen: "Het is 500 liter," en 1 expert zegt: "Ik weet het niet," dan luistert het systeem naar de 8 experts.
• De truc: Als een expert zegt "Ik weet het niet" (terwijl de anderen het wel weten), negeert KohakuRAG die twijfel. Het laat de "zorgzame maar onzekere" expert niet de hele groep domineren.
• Als niemand het weet, zegt het systeem pas echt: "Wij weten het niet." Dit voorkomt dat de computer verzint (hallucineert).

4. De "Nog eens proberen"-knop

Soms zegt de computer: "Ik vind het antwoord niet." In plaats daarvan om de hand te geven, heeft KohakuRAG een retried-mechanisme.
Stel je voor dat je een sleutel in een donkere kamer zoekt. Als je hem niet direct vindt, geef je niet op. Je doet je ogen open, je zoekt met een zaklamp (meer context toevoegen) en je zoekt opnieuw. KohakuRAG doet precies dit: als het eerste antwoord "Ik weet het niet" is, zoekt het nog eens, maar dan met meer informatie. Vaak vindt het het antwoord dan wel.

Waarom won KohakuRAG de wedstrijd?

In de "WattBot 2025 Challenge" moesten systemen vragen beantwoorden over AI-energieverbruik. KohakuRAG won omdat:

1. Het de structuur van de boeken respecteerde (geen losse flarden).
2. Het slimme vertalingen gebruikte om de juiste pagina's te vinden.
3. Het luisterde naar een groep experts in plaats van één persoon.
4. Het niet snel opgaf, maar opnieuw zocht als het antwoord niet direct duidelijk was.

Kortom: KohakuRAG is als een super-georganiseerde bibliothecaris die niet alleen snel zoekt, maar ook weet hoe hij boeken moet lezen, meerdere meningen combineert en nooit iets verzint als hij het niet zeker weet. Hierdoor werd het de nummer 1 in de wereld van AI-vraagbaak-systemen.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "KohakuRAG: A simple RAG framework with hierarchical document indexing" in het Nederlands.

Titel: KohakuRAG: Een eenvoudig RAG-framework met hiërarchische documentindexering

Auteurs: Shih-Ying Yeh, Yueh-Feng Ku, Ko-Wei Huang, Buu-Khang Tu (National Tsing Hua University, Comfy Org Research, Kohaku-Lab)
Context: WattBot 2025 Challenge (een benchmark voor het beantwoorden van technische vragen over AI-energieverbruik met hoge precisie).

---

1. Het Probleem

Retrieval-Augmented Generation (RAG) systemen kampen met fundamentele beperkingen wanneer ze geconfronteerd worden met eisen voor hoge precisie, exacte citaten en numerieke tolerantie (±0,1%):

• Verlies van documentstructuur: Standaard RAG-systemen gebruiken "flat chunking" (vaste lengte segmenten), wat de natuurlijke hiërarchie van documenten (secties, alinea's, zinnen) vernietigt en nauwkeurige citatie-tracking bemoeilijkt.
• Woordenschat-mismatch: Enkele query-vormuleringen missen relevante passages als de terminologie van de gebruiker verschilt van die in de brondocumenten (bijv. "PUE" vs. "power usage effectiveness").
• Stochastische onstabiele antwoorden: Eén enkele inferentie van een LLM levert vaak wisselende antwoorden en citaten op. Systemen neigen ook tot onnodige "abstention" (het weigeren om te antwoorden) zelfs als het bewijs aanwezig is maar moeilijk te lokaliseren.
• Hoge eisen: De WattBot 2025 Challenge vereist dat systemen technische vragen beantwoorden uit 32 documenten (ca. 500K tokens) met exacte bronvermelding en een zeer strakke numerieke tolerantie.

2. Methodologie: KohakuRAG

KohakuRAG lost deze uitdagingen op via een drieledige architectuur die is ontworpen om documentstructuur te behouden, de zoekdekking te vergroten en de stabiliteit van antwoorden te verhogen.

A. Hiërarchische Documentindexering

In plaats van platte chunks, worden documenten geparseerd tot een vier-niveau boomstructuur:

1. Document
2. Sectie
3. Alinea
4. Zin

• Bottom-up Embedding Aggregatie: Embeddings worden berekend voor zinnen en vervolgens geaggregeerd naar bovenliggende niveaus (alinea's, secties) met behulp van een lengte-gewogen gemiddelde. Dit zorgt ervoor dat de semantiek van een sectie de som is van zijn onderdelen, terwijl de structuur behouden blijft.
• Voordeel: Dit biedt natuurlijke grenzen voor citaties op elk niveau en maakt het mogelijk om contextueel relevant bewijs te vinden zonder de structuur te verliezen. Visuele elementen (grafieken/tafels) worden behandeld als speciale alinea-nodes met door een VLM (Vision Language Model) gegenereerde bijschriften.

B. Multi-Query Retrieval met Cross-Query Reranking

Om de woordenschat-kloof te overbruggen:

• Query Planner: Een LLM genereert meerdere semantisch gerelateerde queries per vraag (bijv. het uitbreiden van afkortingen, het herschrijven met synoniemen, het opdelen in sub-vragen).
• Cross-Query Reranking: Resultaten van alle queries worden samengevoegd en opnieuw gerangschikt. Nodes die door meerdere queries worden gevonden, krijgen een hogere rangschikking (consensus-signaal). Dit zorgt ervoor dat de meest relevante passages bovenaan komen, ongeacht de specifieke formulering van de oorspronkelijke vraag.

C. Ensemble Inference met Abstention-Aware Voting

Om de onstabiliteit van LLM-antwoorden te mitigeren:

• Meerdere Runs: Voor elke vraag worden $m$ onafhankelijke inferentie-runs uitgevoerd met een temperatuur > 0.
• Blank Filtering: Als een run "is_blank=true" (onvoldoende bewijs) retourneert, maar er zijn andere runs die een antwoord geven, worden de lege antwoorden genegeerd tijdens het stemmingproces. Dit voorkomt dat conservatieve runs het antwoord domineren wanneer bewijs wel aanwezig is.
• Ensemble Voting: Het uiteindelijke antwoord en de citaten worden bepaald door meerderheidsstemming (majority voting) over de niet-lege antwoorden.

D. Retry Mechanisme

Als het model abstention uitvoert, wordt het systeem geautomatiseerd om de zoekdiepte ($k$) te verhogen en de context opnieuw op te halen. Dit lost vaak fouten op waarbij het bewijs net onder de initiële drempel lag.

3. Belangrijkste Bijdragen

1. Hiërarchische Indexering: Een nieuwe aanpak die documentstructuur behoudt via boomrepresentaties en bottom-up embedding-aggregatie, wat nauwkeurige citatie-tracking mogelijk maakt.
2. Query Planning & Reranking: Een LLM-gedreven planner die vragen uitbreidt naar meerdere formuleringen, gecombineerd met een consensus-gebaseerde reranker.
3. Robuuste Ensemble Inference: Een mechanisme dat meerdere inferenties samenvoegt met specifieke behandeling van "abstention" (weglaten van lege antwoorden), wat de dominante foutmodus (26,8% van de fouten) adresseert.
4. Empirische Validatie: Bewijs dat hiërarchische dense retrieval alleen al concurrerend is (BM25 voegt slechts +3,1% toe), en dat prompt-ordering en retry-mechanismen grotere impact hebben dan hybride zoekstrategieën.

4. Resultaten

De methode werd geëvalueerd op de WattBot 2025 Challenge:

• Leaderboard Prestaties: KohakuRAG behaalde eerste plaats op zowel de publieke als de private leaderboard met een finale score van 0,861.
• Unieke Prestatie: Het was het enige team dat de toppositie behield op beide evaluatiepartities, wat aantoont dat de methode goed generaliseert naar onzichtbare data.
• Ablatie Studies (Impact van componenten):
  • Prompt Ordering: Het plaatsen van de context voor de vraag (in plaats van na) leverde een +80% relatieve verbetering op (verwijzend naar het "lost in the middle" fenomeen).
  • Retry Mechanisme: Leverde +69% verbetering bij lage zoekdiepte door onnodige abstention te corrigeren.
  • Ensemble Voting: Met blank filtering leverde +1,2% absolute verbetering op bij $n=9$.
  • Retrieval: Hiërarchische dense retrieval alleen was zeer sterk; het toevoegen van BM25 (hybride) leverde slechts een marginale winst van +3,1%.

5. Significantie en Conclusie

KohakuRAG demonstreert dat voor complexe, citatie-gebaseerde QA-taken de kwaliteit van de documentstructuur en de robustheid van de inferentie belangrijker zijn dan het simpelweg combineren van verschillende zoekalgoritmen.

• Structuur is cruciaal: Het behoud van de natuurlijke hiërarchie van documenten (via boomstructuren) is essentieel voor het vinden van precieze bronnen.
• Stabiliteit door Ensembles: Het gebruik van ensemble-methoden met slimme filtering van lege antwoorden is effectiever dan het vertrouwen op één enkel model, vooral bij taken met hoge precisie-eisen.
• Open Source: Het framework is beschikbaar gesteld als open-source software, wat bijdraagt aan de gemeenschap voor het bouwen van betrouwbare RAG-systemen.

Het paper concludeert dat de combinatie van hiërarchische indexering, multi-query planning en abstention-aware ensemble voting een nieuwe standaard zet voor RAG-systemen die hoge eisen stellen aan nauwkeurigheid en bronvermelding.