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 (1)
The per-document-type config table is the useful part of this, and it's the step most teams skip — they tune one chunking strategy against a mixed corpus, get a mediocre number, and conclude RAG "doesn't work well" for them. Splitting by document type first and tuning within each is a much better use of the same effort.
On hybrid retrieval: the case that converted me wasn't semantic at all, it was exact-token lookups. Error codes, function names, version strings — pure vector search is genuinely bad at those, because the embedding of
ERR_2047sits near every other error code. BM25 catches them trivially. Worth noting for anyone deciding whether hybrid is worth the extra index.One thing I'd love more detail on: how stable was the Bayesian search result across corpus growth? We tuned retrieval parameters carefully on roughly 50k chunks, then the corpus tripled and the optimum had moved enough that our tuned config was worse than a sane default. We now re-run the sweep on a schedule rather than treating it as a one-time calibration. Did you see drift like that, or did your optimum hold?