KohakuRAG: Climb the Document Tree, Find the Right Evidence

If you've ever built a RAG pipeline at work, you probably know the feeling: retrieval looks fine in demos, then quietly falls apart on real questions.

This post is interesting because it introduces a method that tackles that gap in a very practical, engineering-first way.

The Problems of Standard RAG

Standard RAG is often not reliable enough for hardcore tasks, like the WattBot 2025 Challenge, which demands ±0.1% numeric precision and exact source attribution. In practice, three issues show up again and again:

• Flat chunking breaks semantic boundaries and document structure, which makes precise source tracking much harder and weakens confidence in citation quality.
• A single query can easily miss relevant evidence when the wording does not line up, especially when different terms are used for the same idea.
• A single generation pass is inherently unstable. Both the answer and its citations can vary from run to run, and in some cases, terrified of hallucinating, the system may unnecessarily abstain (outputting a blank) despite the evidence being buried right there in the context.

KohakuRAG: From Flat Chunks to Citation-Ready Trees

At its core, KohakuRAG is a lightweight RAG framework built on a four-level hierarchical document index: document → section → paragraph → sentence.

It preserves document structure through bottom-up embedding aggregation, improves retrieval coverage with LLM-based query planning plus cross-query reranking, and stabilizes final outputs with abstention-aware ensemble voting during generation.

1. Offline hierarchical document indexing

The first step is document parsing.

Each document is organized into a four-level tree structure: document → section → paragraph → sentence. Images and tables are treated as special paragraph nodes, with captions (textual descriptions) generated by Qwen3-VL. Images are extracted with PyMuPDF, then passed to Qwen3-VL for caption generation, with attention placed on technical content, numeric values, axis labels, and similar details.

The document parsing stage is largely rule-based, using layout heuristics to recover section and paragraph structure from PDFs and slides. PyMuPDF handles the basic extraction of text and layout information, and the hierarchy is then identified using rules based on font size, spacing, formatting, and related cues.

• Section detection rules: Text blocks with font sizes at least 1.2 times the median are treated as candidate section titles. Numbering patterns such as 1.2 or A.3 are used as additional signals to confirm section boundaries.
• Paragraph detection rules: Consecutive text blocks with vertical spacing below 1.5 times the line height are merged into a paragraph. Larger spacing is treated as a paragraph boundary.
• Sentence splitting rules: Sentence segmentation is based on NLTK's sent_tokenize, with additional rules tailored for technical writing to handle abbreviations such as Fig., et al., and i.e.

The next step is embedding computation.

Embeddings are aggregated bottom-up using a length-weighted strategy. Leaf nodes (sentences) are embedded directly, while parent nodes are computed from their children using a token-count-weighted average.

KohakuRAG supports two embedding backends: Jina v3 for text retrieval (768-dim, up to 8K context) and Jina v4 for multimodal retrieval, which supports Matryoshka dimensions from 128 to 2048 and can embed both text and images directly. Vector storage is built on sqlite-vec as a lightweight backend.

2. Online Multi-Query Retrieval

The online stage begins with LLM-based query planning.

The original question is expanded into n semantically related queries to cover alternative phrasings, acronym expansion, and sub-query decomposition.

Next comes dense retrieval.

Each planned query is used for cosine similarity search, targeting only sentence and paragraph nodes.

There is also an optional hybrid retrieval step.

BM25-based sparse retrieval can be added as an enhancement, applied only to the paragraph index. Its results are appended directly after the dense retrieval results.

Then comes cross-query reranking.

Three strategies are used here: frequency-based, score-based, and combined reranking. Nodes matched by multiple queries are given higher priority.

Finally, the system performs index traversal and hierarchical context expansion.

For each matched node, the system expands context with parent and sibling nodes. KohakuRAG also describes an optional tree-deduplication step that can remove ancestor–descendant overlap to save context window space, although it was not enabled in the competition submissions.

3. Answering Pipeline

For prompt construction, the retrieved context is placed before the question through a reordered layout. This helps address the lost-in-the-middle issue, while formatted citation markers make attribution easier for the LLM.

For structured output, the LLM is asked to return JSON that includes the answer, the extracted value, citation IDs, an explanation, and a flag indicating whether the result is blank. When the evidence is insufficient, the model is expected to abstain explicitly.

Here, 'citation' is evaluated at the document-ID level rather than as page- or span-level attribution, which is important for interpreting what the system really optimizes.

As shown in Figure 3, a retry mechanism is also introduced. If the LLM returns a blank result, the system increases the retrieval top-k value and runs retrieval again, repeating the process until it gets a valid answer or reaches the retry limit.

4. Multiple Generations, Then Voting

Once the retrieval results are in, the system does not ask the LLM for just a single answer. Instead, it generates multiple answers independently.

At the end, the runs are aggregated through voting. If at least one non-blank answer exists, blank outputs are filtered before voting; if all runs are blank, the system abstains. This makes both answers and citations more stable.

Evaluation

KohakuRAG was developed for the WattBot 2025 Challenge, so the experimental results are mainly based on that benchmark dataset.

There are 32 technical documents on AI energy use, including research papers, technical reports, and presentation slides. The full collection is around 500K tokens and covers topics such as power consumption, water usage, carbon footprint, and performance characteristics of AI models.

On this Challenge, KohakuRAG achieved 0.902 on the public leaderboard and 0.861 on the private leaderboard. It was also the only method to hold the №1 spot on both the public and private leaderboards.

A surprising takeaway: Prompt ordering matters massively. As shown in Figure 6, simply placing the question after the retrieved context (rather than before it) yielded an 80% relative improvement, confirming the "lost in the middle" phenomenon. Sometimes, the biggest gains come from the simplest tweaks.

Another interesting finding: Dense retrieval might be enough. As shown in Figure 7, while hybrid search (Dense + BM25) is commonly treated as a strong baseline, KohakuRAG showed that a well-structured hierarchical dense retrieval alone is highly competitive. Adding BM25 only provided a marginal +3.1% boost.

Thoughts

What makes KohakuRAG valuable is not that any single component feels radically new. The real strength lies in how it tightly integrates four familiar ideas, layered indexing, multi-query retrieval, abstention with retry, and ensemble voting, around a very specific goal: highly accurate answers, precise citations, and the ability to abstain when the evidence is not strong enough.

It also gets at one of the hardest parts of enterprise knowledge QA: traceability matters more than simply getting the answer right. In that sense, the hierarchical tree structure is not just a more elaborate way to chunk documents. It is infrastructure for citations, auditing, and error analysis.

But I have two concerns.

The first concern is that the system is sensitive to document parsing quality, and likely also to the quality of query planning. In real enterprise settings, documents are often full of noisy PDFs and inconsistent layouts. When that happens, rule-based section parsing can go wrong at the foundation level.

The second is cost and latency. Ensemble inference increases cost more or less linearly, and the retry mechanism can make tail latency even worse. So this feels more like a high-reliability, low-throughput design than something naturally suited for high-concurrency online QA.

Reference: KohakuRAG: A simple RAG framework with hierarchical document indexing.