Optimizing RAG at Scale: Chunking, Retrieval, and the Bayesian Search That Cut Latency 40%
How we moved from "semantic search + hope" to a measured, tunable retrieval pipeline with 95% recall@10
The RAG Reality Check
Everyone ships RAG the same way: chunk by 512 tokens, embed with text-embedding-3-small, top-k=5, stuff into context. It works for demos.
Then you hit production:
- Legal contracts: 512 tokens splits clauses mid-sentence
- API docs: 1000-token chunks drown signal in noise
- Customer tickets: Conversational context needs overlap, not fixed windows
- Latency: 500ms embedding + 200ms vector search + 300ms LLM = 1s+ per query
We rebuilt our retrieval layer from first principles. Here's what actually moves metrics.
Chunking: One Size Fits None
# rag/chunking.py
from abc import ABC, abstractmethod
from dataclasses import dataclass
@dataclass
class Chunk:
text: str
metadata: dict
token_count: int
chunk_id: str
class ChunkingStrategy(ABC):
@abstractmethod
def chunk(self, document: str, metadata: dict) -> list[Chunk]: ...
class FixedTokenChunker(ChunkingStrategy):
"""Baseline. Good for homogeneous content."""
def __init__(self, chunk_size=512, overlap=50):
self.chunk_size = chunk_size
self.overlap = overlap
class RecursiveChunker(ChunkingStrategy):
"""Respects structure: markdown headers, code blocks, paragraphs."""
def __init__(self, separators=["\n## ", "\n### ", "\n\n", "\n", " "], chunk_size=512):
self.separators = separators
self.chunk_size = chunk_size
class SemanticChunker(ChunkingStrategy):
"""Uses embedding similarity to find natural boundaries."""
def __init__(self, model="text-embedding-3-small", threshold=0.7):
self.model = model
self.threshold = threshold
class AgenticChunker(ChunkingStrategy):
"""LLM decides boundaries. Expensive but highest quality for complex docs."""
def __init__(self, model="gpt-4o-mini"):
self.model = model
Our production config by document type:
| Document Type | Strategy | Chunk Size | Overlap | Recall@10 |
|---|---|---|---|---|
| Legal contracts | Recursive (clause-aware) | 1024 | 100 | 94% |
| API reference | Recursive (function-aware) | 768 | 50 | 96% |
| Support tickets | Semantic + conversation turns | 512 | 75 | 91% |
| Internal wiki | Agentic (LLM) | 1500 | 200 | 97% |
Hybrid Retrieval: BM25 + Vector + Rerank
Pure vector search misses exact matches (error codes, function names). Pure BM25 misses semantic matches. Hybrid wins.
# rag/retrieval.py
class HybridRetriever:
def __init__(self, vector_store, bm25_index, reranker, weights=(0.4, 0.3, 0.3)):
self.vector = vector_store
self.bm25 = bm25_index
self.reranker = reranker
self.weights = weights # vector, bm25, reranker
async def retrieve(self, query: str, k=20, final_k=5):
# Stage 1: Parallel retrieval
vector_results = await self.vector.search(query, k=k)
bm25_results = await self.bm25.search(query, k=k)
# Stage 2: Reciprocal Rank Fusion
fused = self._rrf(vector_results, bm25_results, k=50)
# Stage 3: Cross-encoder rerank (top 50 → top 5)
reranked = await self.reranker.rerank(query, fused[:50])
return reranked[:final_k]
def _rrf(self, *result_lists, k=60):
"""Reciprocal Rank Fusion — no score calibration needed."""
scores = defaultdict(float)
for results in result_lists:
for rank, doc in enumerate(results):
scores[doc.id] += 1 / (k + rank + 1)
return sorted(scores.items(), key=lambda x: -x[1])
Why cross-encoder rerank? Bi-encoder (embedding) similarity ≈ 0.75 correlation with relevance. Cross-encoder ≈ 0.92. The 50→5 funnel costs 50ms but gains 15% recall.
Query Transformation: Don't Search What User Asked
Users ask badly. Transform first.
# rag/query_transform.py
class QueryTransformer:
def __init__(self, llm_model="gpt-4o-mini"):
self.llm = instructor.from_openai(AsyncOpenAI())
async def expand(self, query: str, context: dict = None) -> list[str]:
"""Generate multiple search queries from one user question."""
class QuerySet(BaseModel):
queries: list[str] = Field(min_length=3, max_length=5)
reasoning: str
result = await self.llm.chat.completions.create(
model=self.model,
response_model=QuerySet,
messages=[
{"role": "system", "content": """
Generate diverse search queries that collectively cover the user's intent.
Include: exact phrasing, synonyms, broader/narrower, hypothetical answer.
"""},
{"role": "user", "content": f"Original: {query}\nContext: {context}"}
],
temperature=0.3,
)
return result.queries
async def decompose(self, query: str) -> list[str]:
"""Break multi-hop questions into sub-questions."""
class SubQuestions(BaseModel):
questions: list[str]
needs_synthesis: bool
return await self.llm.chat.completions.create(
model=self.model,
response_model=SubQuestions,
messages=[...],
)
Query expansion results:
- Single query recall@10: 78%
- 3 expanded queries (union): 94%
- 5 expanded queries (union): 96%
- Cost: 3-5x embedding calls, but parallelizable
Bayesian Optimization: Stop Guessing Hyperparameters
chunk_size=512, top_k=5, similarity_threshold=0.7 — who chose these?
We treat retrieval as a black-box function f(chunk_size, overlap, top_k, weights) → recall@10, latency and optimize with Bayesian search.
# rag/optimization.py
import optuna
from dataclasses import dataclass
@dataclass
class RetrievalConfig:
chunk_size: int
overlap: int
top_k: int
vector_weight: float
bm25_weight: float
rerank_top_k: int
def objective(trial: optuna.Trial) -> tuple[float, float]:
config = RetrievalConfig(
chunk_size=trial.suggest_categorical("chunk_size", [256, 512, 768, 1024, 1536]),
overlap=trial.suggest_int("overlap", 0, 200, step=25),
top_k=trial.suggest_int("top_k", 5, 50, step=5),
vector_weight=trial.suggest_float("vector_weight", 0.1, 0.8),
bm25_weight=trial.suggest_float("bm25_weight", 0.1, 0.8),
rerank_top_k=trial.suggest_int("rerank_top_k", 10, 100, step=10),
)
# Evaluate on golden set (200 queries)
recall, latency = evaluate_config(config, golden_set)
# Multi-objective: maximize recall, minimize latency
return recall, latency / 1000 # seconds
study = optuna.create_study(
directions=["maximize", "minimize"],
sampler=optuna.samplers.TPESampler(multivariate=True),
)
study.optimize(objective, n_trials=100, timeout=3600) # 1 hour
# Pareto frontier gives you the tradeoff curve
pareto = [t for t in study.trials if t.state == TrialState.COMPLETE]
Our Pareto frontier (legal docs, 200-query golden set):
| Config | Recall@10 | Latency (p95) | Use Case |
|---|---|---|---|
| Conservative | 91% | 180ms | High-throughput API |
| Balanced (prod) | 95% | 320ms | Default |
| Aggressive | 97% | 580ms | High-stakes legal/medical |
Production Metrics Dashboard
# rag/metrics.py
from prometheus_client import Histogram, Counter, Gauge
RETRIEVAL_LATENCY = Histogram("rag_retrieval_latency_seconds", "End-to-end retrieval time")
RECALL_AT_K = Gauge("rag_recall_at_k", "Recall@k on golden set", ["k"])
QUERY_EXPANSION_COUNT = Counter("rag_query_expansions_total", "Number of expanded queries")
RERANKER_LATENCY = Histogram("rag_reranker_latency_seconds", "Cross-encoder rerank time")
class InstrumentedRetriever(HybridRetriever):
async def retrieve(self, query, k=20, final_k=5):
with RETRIEVAL_LATENCY.time():
expanded = await self.transformer.expand(query)
QUERY_EXPANSION_COUNT.inc(len(expanded))
results = await super().retrieve(expanded, k, final_k)
# Track recall on sampled golden queries (1% of traffic)
if random.random() < 0.01:
RECALL_AT_K.labels(k=10).set(self._eval_recall(query, results))
return results
Results: 6 Months of Iteration
| Metric | Baseline (naive) | Optimized | Improvement |
|---|---|---|---|
| Recall@10 | 78% | 95% | +17 pp |
| Latency p95 | 850ms | 320ms | -62% |
| Hallucination rate | 12% | 3% | -75% |
| Cost/query | $0.008 | $0.005 | -38% |
The Checklist for Your RAG
- [ ] Chunk by document structure, not fixed tokens
- [ ] Hybrid retrieval (BM25 + vector + rerank) — never single modality
- [ ] Query expansion for ambiguous/short queries
- [ ] Golden dataset with stratified cases (version it in Git)
- [ ] Bayesian optimization of hyperparams (re-run monthly)
- [ ] Instrumentation on every retrieval (latency, recall sampling)
- [ ] A/B framework for retrieval changes (feature flags)
The Mental Shift
Retrieval is infrastructure, not afterthought.
- Treat chunking strategies as first-class code (versioned, tested, reviewed)
- Golden dataset = your most valuable IP (curate it religiously)
- Every retrieval change = eval run (enforced by CI)
- Regression alerts = paging alerts (not email digests)
Your users don't care about your embedding model. They care that the answer is right. Automated evaluation is how you guarantee that at scale.
Code: github.com/yourname/rag-eval-framework |
Discussion: Hacker News |
Follow: @yourname
Top comments (0)