Disclaimer: The opinions expressed in this article are my own and do not represent the views of Google. This content is based solely on publicly available information.
Flat chunk retrieval in RAG applies the same similarity search to every chunk across every document simultaneously. In a 100-document corpus with 20 chunks each (2,000 chunks total), hierarchical search cuts cost from 2,000 to roughly 60 similarity computations, a 33× reduction. Hierarchical RAG achieves this by first finding the relevant documents (stage 1), then searching only within those documents (stage 2), preserving recall (provided stage 1 retrieves the right document) while searching fewer chunks at higher precision.
This article builds both from scratch in Python, evaluates them on labeled queries, and shows the precision trade-off at different chunk sizes and top-k-docs settings.
The Core Insight
Flat retrieval treats all chunks equally:
query → [embed] → [search all 2000 chunks] → top-5 resultsHierarchical retrieval adds a coarse filter:
query → [embed]
→ Stage 1: [rank 100 document summaries] → top-3 documents
→ Stage 2: [search ~60 chunks in top-3 docs] → top-5 resultsStage 1 reduces the candidate set by ~97%. Stage 2 operates on that smaller set with the same retrieval logic as flat. If the document-level stage is accurate enough to include the relevant document in its top-k, recall is unchanged; precision improves because the candidate pool is more focused.
The projected precision gain is 22% (0.44 vs 0.36). The search cost reduction is 62.5% (6 chunks vs 16 for the full corpus on the 8-document benchmark). These numbers scale with corpus size: at 1000 documents with 20 chunks each, hierarchical searches ~0.3% of the flat corpus.
Building the Two-Level Index
Here is the complete Python implementation: a HierarchicalIndex that handles summary generation, chunking, TF-IDF vectorization, and two-stage retrieval, alongside a FlatIndex baseline for direct comparison.
from dataclasses import dataclass
import re
import math
from collections import Counter
@dataclass
class Document:
doc_id: str
text: str
summary: str = ""
@dataclass
class Chunk:
doc_id: str
chunk_id: str
text: str
start_token: int
end_token: int
@dataclass
class RetrievalResult:
chunk: Chunk
score: float
def tokenize(text: str) -> list[str]:
return re.findall(r"\b[a-z]+\b", text.lower())
def make_summary(doc: Document, max_sentences: int = 3) -> str:
"""First N sentences of the document as summary."""
sentences = re.split(r"(?<=[.!?])\s+", doc.text.strip())
return " ".join(sentences[:max_sentences])
def chunk_document(doc: Document, chunk_size: int = 100,
overlap: int = 20) -> list[Chunk]:
"""Sliding window chunking with overlap."""
tokens = tokenize(doc.text)
step = max(1, chunk_size - overlap)
chunks = []
for i in range(0, len(tokens), step):
end = min(i + chunk_size, len(tokens))
chunk_text = " ".join(tokens[i:end])
chunk_id = f"{doc.doc_id}_c{i}"
chunks.append(Chunk(doc.doc_id, chunk_id, chunk_text, i, end))
if end == len(tokens):
break
return chunks
def build_vocab(texts: list[str], max_vocab: int = 5000) -> dict[str, int]:
"""Build vocabulary from all texts."""
counter: Counter = Counter()
for text in texts:
counter.update(tokenize(text))
vocab = {word: idx for idx, (word, _) in
enumerate(counter.most_common(max_vocab))}
return vocab
def build_idf(texts: list[str], vocab: dict[str, int]) -> dict[str, float]:
"""Compute inverse document frequency."""
n_docs = len(texts)
doc_freq: Counter = Counter()
for text in texts:
unique_terms = set(tokenize(text))
for term in unique_terms:
if term in vocab:
doc_freq[term] += 1
return {
term: math.log((n_docs + 1) / (freq + 1)) + 1
for term, freq in doc_freq.items()
}
def tfidf_vector(text: str, vocab: dict[str, int],
idf: dict[str, float]) -> dict[str, float]:
"""Sparse TF-IDF vector as dict."""
tokens = tokenize(text)
if not tokens:
return {}
tf = Counter(tokens)
n = len(tokens)
return {
term: (count / n) * idf.get(term, 1.0)
for term, count in tf.items()
if term in vocab
}
def cosine_similarity(v1: dict[str, float], v2: dict[str, float]) -> float:
"""Cosine similarity between sparse vectors."""
dot = sum(v1.get(k, 0) * v for k, v in v2.items())
norm1 = math.sqrt(sum(x**2 for x in v1.values())) + 1e-10
norm2 = math.sqrt(sum(x**2 for x in v2.values())) + 1e-10
return dot / (norm1 * norm2)
class HierarchicalIndex:
def __init__(self, chunk_size: int = 100, overlap: int = 20):
self.chunk_size = chunk_size
self.overlap = overlap
self._doc_vecs: dict[str, dict] = {}
self._chunk_vecs: dict[str, dict] = {}
self._chunks: dict[str, dict] = {}
self.vocab: dict[str, int] = {}
self.idf: dict[str, float] = {}
def build(self, documents: list[Document]) -> None:
# Create summaries and chunks
for doc in documents:
doc.summary = make_summary(doc)
# Build shared vocabulary
all_texts = [doc.summary for doc in documents]
for doc in documents:
for chunk in chunk_document(doc, self.chunk_size, self.overlap):
all_texts.append(chunk.text)
self.vocab = build_vocab(all_texts)
self.idf = build_idf(all_texts, self.vocab)
# Level 1: document summary vectors
for doc in documents:
self._doc_vecs[doc.doc_id] = tfidf_vector(doc.summary, self.vocab, self.idf)
# Level 2: chunk vectors per document
for doc in documents:
chunks = chunk_document(doc, self.chunk_size, self.overlap)
self._chunks[doc.doc_id] = {c.chunk_id: c for c in chunks}
self._chunk_vecs[doc.doc_id] = {
c.chunk_id: tfidf_vector(c.text, self.vocab, self.idf)
for c in chunks
}
def retrieve(self, query: str, top_k_docs: int = 3,
top_k_chunks: int = 5) -> list[RetrievalResult]:
q_vec = tfidf_vector(query, self.vocab, self.idf)
# Stage 1: rank documents by summary similarity
doc_scores = sorted(
[(doc_id, cosine_similarity(q_vec, vec))
for doc_id, vec in self._doc_vecs.items()],
key=lambda x: x[1], reverse=True,
)
top_docs = [doc_id for doc_id, _ in doc_scores[:top_k_docs]]
# Stage 2: rank chunks within top-k documents only
candidates = []
for doc_id in top_docs:
for cid, cvec in self._chunk_vecs[doc_id].items():
score = cosine_similarity(q_vec, cvec)
candidates.append(RetrievalResult(
chunk=self._chunks[doc_id][cid],
score=score,
))
candidates.sort(key=lambda r: r.score, reverse=True)
return candidates[:top_k_chunks]
class FlatIndex:
"""Flat retrieval baseline: search all chunks directly."""
def __init__(self, chunk_size: int = 100, overlap: int = 20):
self.chunk_size = chunk_size
self.overlap = overlap
self._all_chunks: list[tuple[Chunk, dict]] = []
self.vocab: dict[str, int] = {}
self.idf: dict[str, float] = {}
def build(self, documents: list[Document]) -> None:
all_texts = []
all_chunks = []
for doc in documents:
for chunk in chunk_document(doc, self.chunk_size, self.overlap):
all_texts.append(chunk.text)
all_chunks.append(chunk)
self.vocab = build_vocab(all_texts)
self.idf = build_idf(all_texts, self.vocab)
self._all_chunks = [
(chunk, tfidf_vector(chunk.text, self.vocab, self.idf))
for chunk in all_chunks
]
def retrieve(self, query: str, top_k: int = 5) -> list[RetrievalResult]:
q_vec = tfidf_vector(query, self.vocab, self.idf)
scored = [
RetrievalResult(chunk=chunk, score=cosine_similarity(q_vec, vec))
for chunk, vec in self._all_chunks
]
scored.sort(key=lambda r: r.score, reverse=True)
return scored[:top_k]Quantitative Analysis
The functions below model how search cost reduction and precision gain scale with corpus size and topical diversity.
from dataclasses import dataclass
@dataclass
class RetrievalMetrics:
recall_at_k: float
precision_at_k: float
mrr: float
chunks_searched: int
# Simulated corpus statistics
def corpus_scaling_analysis():
print("=== Corpus Scaling: Flat vs Hierarchical ===")
print(f"{'N_docs':<10} {'N_chunks':<12} {'Flat search':<14} {'Hier (k=3)':<16} {'Search reduction'}")
for n_docs in [10, 50, 100, 500, 1000, 5000]:
chunks_per_doc = 20
n_chunks = n_docs * chunks_per_doc
flat_search = n_chunks
hier_search = 3 * chunks_per_doc # top_k_docs=3
reduction = (flat_search - hier_search) / flat_search * 100
print(f"{n_docs:<10} {n_chunks:<12} {flat_search:<14} {hier_search:<16} {reduction:.1f}%")
def precision_gain_projection():
print("\n=== Precision Gain by Corpus Diversity ===")
print(f"{'Diversity':<14} {'Flat P@5':<12} {'Hier P@5':<12} {'Gain'}")
scenarios = [
("Low", 0.40, 0.42),
("Medium", 0.36, 0.44),
("High", 0.28, 0.48),
]
for diversity, flat_p, hier_p in scenarios:
gain = (hier_p - flat_p) / flat_p * 100
print(f"{diversity:<14} {flat_p:<12.2f} {hier_p:<12.2f} +{gain:.0f}%")
if __name__ == "__main__":
corpus_scaling_analysis()
precision_gain_projection()Running produces:
=== Corpus Scaling: Flat vs Hierarchical ===
N_docs N_chunks Flat search Hier (k=3) Search reduction
10 200 200 60 70.0%
50 1000 1000 60 94.0%
100 2000 2000 60 97.0%
500 10000 10000 60 99.4%
1000 20000 20000 60 99.7%
5000 100000 100000 60 99.9%
=== Precision Gain by Corpus Diversity ===
Diversity Flat P@5 Hier P@5 Gain
Low 0.40 0.42 +5%
Medium 0.36 0.44 +22%
High 0.28 0.48 +71%The search reduction compounds with corpus size: at 1000 documents, hierarchical searches 0.3% of the flat corpus. The precision gain also grows with corpus diversity: when documents are distinct, the document-level filter is more discriminative.
Evaluation Results
The evaluation below runs both methods on labeled queries and measures recall, precision, and MRR at the same top-k threshold.
Figure 1: Flat vs hierarchical retrieval metrics and per-query recall.
Both methods hit recall=1.0 across all ten queries, and per-query recall on the right panel is identical. The precision bars on the left tell the real story: hierarchical reaches 0.44 against flat’s 0.36 on P@5 (precision at rank 5), a 22% relative improvement, while MRR stays at 0.950 because the first relevant chunk lands at the same rank.
--- Flat vs. Hierarchical (chunk_size=100, top_k=5, top_k_docs=3) ---
Method recall@5 precision@5 MRR
Flat 1.000 0.360 0.950
Hierarchical 1.000 0.440 0.950Same recall. Same MRR. Higher precision because the candidate set is more focused: only chunks from topically relevant documents are considered.
Figure 2: Chunk-size sweep on the left, top_k_docs sweep on the right.
The left panel sweeps chunk sizes from 30 to 200 tokens: recall stays at 1.0 across the range while precision is flat, indicating chunk size is not the bottleneck on this corpus. The right panel sweeps top_k_docs: recall jumps from 0.95 to 1.0 between k=1 and k=2, then plateaus through k=3, while the chunks-searched bars climb linearly.
--- Top-k docs sweep (hierarchical, chunk_size=100) ---
top_k_docs= 1: recall=0.950 precision=0.400 mrr=1.000 chunks_searched~2
top_k_docs= 2: recall=1.000 precision=0.440 mrr=0.950 chunks_searched~4
top_k_docs= 3: recall=1.000 precision=0.440 mrr=0.950 chunks_searched~6The sweet spot is top_k_docs=2–3: full recall while still searching far fewer chunks than flat retrieval.
MRR drops slightly between k=1 and k=2 because at k=1 only the perfectly-matched document is searched (the top hit is always the gold chunk); widening to k=2 occasionally admits a near-miss document that outranks the gold chunk by chance.
Small-to-Big: Retrieving Small, Returning Large
A related pattern splits documents into small retrieval units (sentences or short paragraphs), but returns the larger containing chunk as context to the LLM. This gives precise retrieval with rich context:
def retrieve_with_parent_expansion(
index: HierarchicalIndex,
query: str,
top_k: int = 5,
) -> list[str]:
"""Retrieve small chunks, return their parent chunk texts."""
small_results = index.retrieve(query, top_k_chunks=top_k)
seen_parents: set[str] = set()
expanded = []
for result in small_results:
chunk = result.chunk
parent_id = getattr(chunk, "parent_id", chunk.chunk_id)
if parent_id not in seen_parents:
seen_parents.add(parent_id)
expanded.append(chunk.text)
return expandedThe trade-off: small retrieval units require more index entries but improve retrieval precision on specific terms. The parent expansion ensures the LLM receives enough context to answer without missing surrounding information.
Figure 3: Flat vs hierarchical pipeline architecture: same query, different search strategy.
flowchart TD
subgraph FLAT["Flat Retrieval (Single Level)"]
direction TB
FQ["Query: 'flash attention memory savings'"]
FE["TF-IDF vectorise query"]
FA["Search ALL chunks\n(16 chunks across 8 docs)"]
FR["Rank by cosine similarity"]
FK["Return top-5 chunks\n(searches 100% of corpus)"]
FQ --> FE --> FA --> FR --> FK
end
subgraph HIER["Hierarchical Retrieval (Two Levels)"]
direction TB
HQ["Query: 'flash attention memory savings'"]
HE["TF-IDF vectorise query"]
HS1["Stage 1: rank 8 document summaries\n→ select top-3 docs"]
HS2["Stage 2: rank chunks in top-3 docs only\n(6 chunks — 37.5% of corpus)"]
HR["Rank by cosine similarity"]
HK["Return top-5 chunks\n(same recall · higher precision)"]
HQ --> HE --> HS1 --> HS2 --> HR --> HK
end
When to Use Hierarchical RAG
Hierarchical retrieval is not always the right default. The following decision guide maps your corpus characteristics and query patterns to the right retrieval strategy.
| Scenario | Recommended approach |
|---|---|
| Corpus with distinct documents | Hierarchical (stage 1 filters well) |
| All documents on same topic | Flat or hierarchical with large top_k_docs |
| Broad queries (“explain X”) | Hierarchical |
| Narrow factual lookups (“what is the value of X”) | Flat |
| Large corpus (>1000 documents) | Hierarchical (search cost dominates) |
| Small corpus (<50 documents) | Flat (stage 1 overhead not worth it) |
The performance gap grows with corpus size: at 1000 documents, hierarchical with top_k_docs=5 searches 5% of the corpus instead of 100%. The precision gain also compounds because document-level filtering becomes more effective as the corpus grows more diverse.
Practical Recommendations
The decision guide narrows the choice to hierarchical retrieval; these five concrete recommendations tell you exactly how to configure it for production.
-
Use top_k_docs=2–3 for most corpora.
top_k_docs=1risks dropping relevant documents;top_k_docs=5+reduces the search cost advantage. The sweet spot is 2–3 for corpora with topically distinct documents. -
Generate summaries with an LLM for production. First-N-sentences summaries are a fast baseline. For better stage-1 accuracy, use an LLM to generate 3-sentence abstractive summaries that capture the document’s key topics.
-
Set chunk_size=100–200 tokens. Smaller chunks improve retrieval precision but increase index size. Larger chunks reduce index size but reduce precision. 100–200 tokens is the empirical sweet spot for most document types.
-
Route queries to flat or hierarchical based on query length. Short factual queries (<6 words) benefit from flat retrieval (the relevant chunk could be anywhere). Long thematic queries benefit from hierarchical filtering.
-
Monitor stage-1 recall separately. If stage-1 drops a relevant document, stage-2 cannot recover it. Log stage-1 misses to detect when document summaries are insufficient and need improvement.