KohakuRAG: A simple RAG framework with hierarchical document indexing

Shih-Ying Yeh; Yueh-Feng Ku; Ko-Wei Huang; Buu-Khang Tu

cs.CLMar 8, 2026

KohakuRAG: A simple RAG framework with hierarchical document indexing

Authors: Shih-Ying Yeh, Yueh-Feng Ku, Ko-Wei Huang, Buu-Khang Tu

Organizations: National Tsing Hua University · Comfy Org Research · Kohaku-Lab

Abstract

Retrieval-augmented generation (RAG) systems that answer questions from document collections face compounding difficulties when high-precision citations are required: flat chunking strategies sacrifice document structure, single-query formulations miss relevant passages through vocabulary mismatch, and single-pass inference produces stochastic answers that vary in both content and citation selection. We present KohakuRAG, a hierarchical RAG framework that preserves document structure through a four-level tree representation (document $\rightarrow$ section $\rightarrow$ paragraph $\rightarrow$ sentence) with bottom-up embedding aggregation, improves retrieval coverage through an LLM-powered query planner with cross-query reranking, and stabilizes answers through ensemble inference with abstention-aware voting. We evaluate on the WattBot 2025 Challenge, a benchmark requiring systems to answer technical questions from 32 documents with $\pm$ 0.1% numeric tolerance and exact source attribution. KohakuRAG achieves first place on both public and private leaderboards (final score 0.861), as the only team to maintain the top position across both evaluation partitions. Ablation studies reveal that prompt ordering (+80% relative), retry mechanisms (+69%), and ensemble voting with blank filtering (+1.2pp) each contribute substantially, while hierarchical dense retrieval alone matches hybrid sparse-dense approaches (BM25 adds only +3.1pp). We release KohakuRAG as open-source software at https://github.com/KohakuBlueleaf/KohakuRAG.

Figures & tables

Figure 1: Overview of KohakuRAG. Left (Hierarchical Indexing) : Documents are parsed into tree structures with sections, paragraphs (Para), and sentences (S). Sentence embeddings are computed and aggregated bottom-up to parent levels, then stored in a Vector DB. Center (Multi-Query Retrieval) : Given a question, the Query Planner (LLM) generates multiple related queries, each retrieving Top-K results that are merged via Cross-Query Reranking. Right (Ensemble Inference) : Context and question are sent to the LLM for $m$ independent runs; blank responses are filtered (X), and majority voting produces the final answer.

Figure 2: Hierarchical context expansion. Given a query, we retrieve Top-K nodes across multiple granularity levels (sentences in cyan, paragraphs in green, sections in purple). Section-level nodes are filtered out as too coarse (marked with X). The remaining nodes are expanded by adding their parent nodes to provide broader context. For example, a matched sentence sec1:p1:s1 brings its parent paragraph sec1:p1 , while a matched paragraph sec2:p2 brings its parent section sec2 .

Figure 3: Retry mechanism flow. When the LLM outputs is_blank=true (insufficient evidence), the system increases top- $k$ and re-retrieves context. This loop continues until a valid answer is produced or the retry limit is reached, at which point the system outputs a blank response.

Figure 4: Ensemble voting with blank handling. Case 1 (Mixed Responses) : When non-blank answers exist among 5 runs {a, a, blank, b, a}, blank responses are filtered (X) before majority voting, yielding {a, a, b, a} with majority “a” (3 votes). Case 2 (All Blank) : When all runs return blank, the system correctly abstains rather than hallucinating, preserving the hallucination score.

Figure 5: Dual-path image processing. Left : Caption-based retrieval via Qwen3-VL captioning and Jina v3/v4 text embedding. Right : Direct image embedding via Jina v4’s multimodal encoder.

Figure 6: UMAP visualization of Jina v4 embeddings for all indexed nodes, colored by source document. Points are differentiated by node type (documents, sections, sentences, paragraph averages, images, and full paragraphs). We annotate three types of regions: green circles mark document-dominant clusters where one paper dominates 50–99% of nodes with minor contributions from related documents; red circles highlight cross-document semantic clusters where multiple documents mix, indicating shared topics (e.g., energy measurement methodologies); blue circles show image-text alignment regions where image embeddings (stars) cluster with semantically related sentences. This validates that our hierarchical embeddings preserve document coherence while enabling cross-document similarity matching.

Component	Setting	$k$ =4	$k$ =8	$k$ =16	Avg Rank
Embedding	Jina v3 (text only)	0.752 ∗∗	0.801 †	0.824 ∗	3.00
	Jina v3 (+ captions)	0.859	0.836 †	0.824 ∗∗	2.00
	Jina v4 (+ images)	0.869	0.868	0.890	1.00
Planner	2 queries	0.562 ∗∗	0.466 ∗	0.558 ∗∗	3.00
	4 queries	0.678	0.830	0.811	1.67
	6 queries	0.787	0.809	0.838	1.33

Table 1: Retrieval component ablations (final score). Best per column in bold , second best underlined . Significance vs. best: †

p

<

0.10, ∗

p

<

0.05, ∗∗

p

<

0.01.

Component	Setting	$k$ =4	$k$ =8	$k$ =16	Avg Rank
Prompt Order	Standard (Q $\rightarrow$ C)	0.418 ∗	0.485 ∗∗∗	0.704	2.00
Prompt Order	Reordered (C $\rightarrow$ Q)	0.752	0.783	0.811	1.00
BM25 top- $k$	0 (dense only)	0.688 ∗	0.741	0.788	3.67
	2	0.701	0.822	0.812	2.00
	4	0.763	0.801	0.819	1.33
	8	0.720	0.720 ∗∗	0.806	3.00

Table 2: Context enhancement ablations (final score). Best per column in bold , second best underlined . Significance vs. best: †

p

<

0.10, ∗

p

<

0.05, ∗∗

p

<

0.01, ∗∗∗

p

<

0.001.

Model	$k$ =4	$k$ =8	$k$ =16	Rank
GPT-5-nano	.783 ∗∗	.746 ∗∗	.787 ∗∗	7.00
GPT-5-mini	.812 ∗	.804 ∗	.847	4.67
GPT-oss-120B	.805 †	.763 ∗	.833 ∗	5.67
Mistral-large	.853 ∗	.862	.816 †	4.67
Gemini-3-pro	.871 †	.873	.875	2.67
Kimi-k2	.872 ∗	.889	.874	2.00

Table 3: LLM comparison (final score). Best in bold , second underlined . Significance: †

p

<

0.10, ∗

p

<

0.05, ∗∗

p

<

0.01.

ignore_blank	$n$ =1	$n$ =5	$n$ =9	$n$ =13	$n$ =15
Off (1.88)	.799	.866 ∗∗	.882 ∗∗	.883 ∗∗∗	.884 ∗∗
On ( 1.12 )	.799	.886	.894	.896	.895
$\Delta$	–	+.020	+.012	+.013	+.011

Table 4: Ensemble analysis. Left: ignore_blank ablation (avg rank in parentheses); all differences for

n

\geq

3 are significant (

p

<

0.01). Right: voting strategy comparison at

n

=15; differences among top strategies are not significant (

p

>

0.1).

Model	Config	$n$	Pub	Priv
GPT-oss-120B	No img/BM25	5	.884	.852
GPT-oss-120B	+BM25	7	.862	.857
Gemini-3-pro	+Images	1	.901	.839
Top-5 mixed ensemble		–	.902	.861

Table 5: WattBot 2025 test set results and leaderboard comparison. Left : Our submitted configurations (all use Jina v4 with hierarchical indexing). Right : Public vs. private score dynamics across top teams. Only KohakuRAG maintained #1 on both partitions.

Appendix figures & tables25 assets

Supplementary material from the paper’s appendix.

Appendix

Mode	Matches	Snippets	Context (chars)
No dedup	32.0 $\pm$ 0.0	64.0 $\pm$ 0.0	148,125 $\pm$ 34,415
Node-ID dedup	20.9 $\pm$ 3.9	41.8 $\pm$ 7.8	96,332 $\pm$ 24,577
Tree dedup	15.7 $\pm$ 4.4	31.4 $\pm$ 8.9	71,956 $\pm$ 23,449

Appendix

Table 6: Context statistics under different deduplication modes (mean

\pm

std over 41 training questions, 4 queries

\times

16 retrievals each). Node-ID dedup removes cross-query duplicates (

-

35.0%); tree dedup additionally removes hierarchical overlap (

-

51.4% cumulative).

Parameter	Values Tested	Default	Notes
top- $k$ (per query)	4, 8, 16	16	Higher = more coverage
top- $k_{\text{final}}$	16, 32, None	32	Post-rerank truncation
Planner queries	2, 4, 6	4	More = better coverage
Max retries	0, 1, 2	0	On abstention
BM25 top- $k$	0, 2, 4, 8	0	0 = dense only
Ensemble size	1, 3, 5, 7, 9, 11, 13, 15	1	Odd numbers only

Appendix

Table 7: Hyperparameter search space and defaults.

Model	Jina v3 text	Jina v3 +img	Jina v4 +img
GPT-5-nano	0.825	0.798	0.807
GPT-5-mini	0.832	0.879	0.841
GPT-oss-120B	0.436	0.501	0.431
Mistral-large-2512	0.844	0.856	0.864
Kimi-k2-thinking	0.867	0.879	0.872

Appendix

Table 8: LLM

\times

Embedding configuration (final score).

Figure 9: Send images to LLM $\times$ Image top- $k$ configuration. Sending 2 images provides optimal performance.

Send to LLM	img $k$ =0	img $k$ =2	img $k$ =4	img $k$ =8
False (caption only)	0.875	0.875	0.876	0.881
True (with images)	0.875	0.887	0.877	0.853

Appendix

Table 9: Image retrieval configuration with Grok-4.1-fast (final score).

Figure 16: LLM model $\times$ Top- $k$ configuration.

Strategy	$n$ =1	$n$ =3	$n$ =5	$n$ =7	$n$ =9	$n$ =11	$n$ =13	$n$ =15
Independent	.799	.841	.869	.876	.886	.883	.888	.889
RefPriority	.799	.836	.864	.871	.883	.880	.886	.889
AnswerPriority	.799	.837	.868	.876	.886	.884	.889	.892
Union	.799	.836	.866	.872	.882	.878	.883	.884
Intersection	.799	.838	.868	.875	.886	.882	.888	.890

Appendix

Table 10: Ensemble size vs. voting strategy (final score, mean over 32 subsamples).

ignore_blank	$n$ =1	$n$ =3	$n$ =5	$n$ =7	$n$ =9	$n$ =11	$n$ =13	$n$ =15
Off	.799	.836	.866	.872	.882	.878	.883	.884
On	.799	.860	.886	.886	.894	.891	.896	.895
$\Delta$	–	+.024	+.020	+.014	+.012	+.013	+.013	+.011

Appendix

Table 11: Effect of ignore_blank on ensemble performance (union voting).

Metric	Value
Total predictions analyzed	2,583
Correct predictions	1,942 (75.2%)
Total errors	641 (24.8%)
LLM configurations tested	7
Top- $k$ settings tested	3 (4, 8, 16)
Random seeds per config	3

Appendix

Table 12: Summary of error analysis across all sweep configurations.

Figure 19: Distribution of error types across 641 errors from 63 sweep runs. The top three categories (abstention, reference mismatch, value selection) account for nearly three-quarters of all errors.

Figure 20: Error rates by LLM model and retrieval depth. Grok-4.1-fast achieves the lowest error rates across all configurations, while smaller models (GPT-5-nano) show consistently higher error rates.

Panel of Figure 23 — Figure 23: WattBot 2025 Challenge leaderboards. (a) Public leaderboard visible during the competition. (b) Final private leaderboard revealed after the submission deadline. KohakuRAG maintained rank #1 across both partitions.

Metric	Value
Total teams	66
Mean $\Delta$ (private $-$ public)	-0.0106
Standard deviation of $\Delta$	0.0452
Median $\Delta$	-0.0160
Minimum $\Delta$	-0.2130
Maximum $\Delta$	0.0980

Appendix

Table 13: Summary statistics for public-to-private score transitions across all participating teams. The score delta (

\Delta

) is computed as private score minus public score, where negative values indicate performance degradation on the private partition.

Figure 24: Relationship between public score and score delta. Each point represents a team; KohakuRAG is marked with a red star. The negative trend reflects variance between the two test distributions rather than indicating that high public scores are inherently problematic.

Figure 25: Distribution of score deltas (private $-$ public) across all teams. KohakuRAG’s delta is marked for reference.

Pub Rank	Pub Score	Priv Score	Priv Rank	Score $\Delta$	Rank $\Delta$
#1 $\star$	0.902	0.861	#1	-0.041	0
#2	0.886	0.840	#4	-0.046	-2
#3	0.868	0.831	#9	-0.037	-6
#4	0.862	0.831	#8	-0.031	-4
#5	0.856	0.806	#17	-0.050	-12
#6	0.856	0.839	#6	-0.017	0

Appendix

Table 14: Public-to-private transition for the top 10 public leaderboard teams. The

\Delta

column shows score change, and Rank

\Delta

shows position change (positive = improved). KohakuRAG (

\star

) was the only team to maintain first place across both partitions.

Figure 26: Rank transitions from public to private leaderboard for top 20 public teams. Lines connect each team’s public rank (left) to private rank (right). KohakuRAG (red) maintained rank #1; green lines indicate teams that improved, gray lines indicate teams that dropped.

Rank Band	$n$	Mean $\Delta$	Std $\Delta$	Median $\Delta$
1–10	10	-0.0323	0.0132	-0.0365
11–20	10	-0.0241	0.0289	-0.0205
21–30	10	+0.0016	0.0284	-0.0040
31–50	20	-0.0181	0.0597	-0.0190
51–100	16	+0.0131	0.0414	+0.0230

Appendix

Table 15: Score delta statistics stratified by public rank band. Top-ranked teams show more negative mean deltas, while mid-to-low ranked teams show near-zero or positive mean deltas.

Public	Private	$\Delta$ Rank	$\Delta$ Score
#21	#2	+19	+0.055
#20	#3	+17	+0.044
#25	#15	+10	+0.035
#23	#16	+7	+0.024
#10	#5	+5	-0.007

Appendix

Table 16: Teams with the largest rank improvements (left) and largest rank drops (right). Notable gainers came from outside the top 20 public, while notable losers fell from top-10 positions.

Retrieval-Augmented Generation

Retrieval-augmented generation connects language models to external knowledge so answers can use current, domain-specific evidence.

See all 766 papers

RIT-RAG: Navigating Document Corpora with Retrieval-Induced Trees
Oct 8, 2026Meghanadh Pulivarthi, Swaraj Kumar Biswal, Kushagra Bhushan +3Retrieval-Augmented Generation Agentic RAG
IBM · IIT Kharagpur
SproutRAG: Attention-Guided Tree Search with Progressive Embeddings for Long-Document RAG
Jun 16, 2026Amirhossein Abaskohi, Issam H. Laradji, Peter West +1Document Chunking Retrieval-Augmented Generation
University of British Columbia · ServiceNow Research
StratRAG: A Multi-Hop Retrieval Evaluation Dataset for Retrieval-Augmented Generation Systems
Mar 6, 2026Aryan PatodiyaRAG Evaluation Retrieval-Augmented Reasoning
Department of Computer Science, California State University, Fresno, Fresno, CA, USA

Oct 8, 2026J/K move · Enter open · S save

RIT-RAG: Navigating Document Corpora with Retrieval-Induced Trees

Meghanadh Pulivarthi, Swaraj Kumar Biswal, Kushagra Bhushan, Yatin Nandwani+2

IBM · IIT Kharagpur

Retrieval-Augmented Generation Agentic RAG Question Answering

Open paper AskPDF arXiv

Retrieval-augmented generation (RAG) grounds language models in external corpora. Agentic RAG enables iterative search, yet exposes the model to isolated chunks without document structure, making it difficult to distinguish relevant evidence from chunks that merely resemble the query. Structure-aware methods such as PageIndex navigate document structure but cannot scale to the structures of large corpora, which do not fit in the LLM context. Hence, they first commit to a single document using a document retriever and cannot recover from a wrong choice. We propose RIT-RAG (Retrieval-Induced Tree RAG), which combines content retrieval with structural navigation. Offline, RIT-RAG builds a tree for each document from its table of contents or sitemap. At query time, it retrieves a broad set of chunks and uses their positions to induce manageable sub-trees, potentially across multiple documents. An LLM agent navigates these sub-trees, selectively reads promising nodes, and reformulates queries when needed. Thus, retrieval proposes where to look, while the agent decides what to read. Across financial, scientific, and customer-support benchmarks, RIT-RAG achieves the highest answer accuracy among vanilla, graph-based, and agentic baselines. On EntQABench, our new benchmark of 2.84 million technical-documentation webpages, it improves accuracy by 6.8 to 11.4 points over the strongest baseline across three LLMs.

Abstract

Retrieval-Augmented Generation

Explore similar work

RIT-RAG: Navigating Document Corpora with Retrieval-Induced Trees

SproutRAG: Attention-Guided Tree Search with Progressive Embeddings for Long-Document RAG

StratRAG: A Multi-Hop Retrieval Evaluation Dataset for Retrieval-Augmented Generation Systems