DEV Community

Imus
Imus

Posted on

Optimizing RAG at Scale: Chunking, Retrieval, and the Bayesian Search That Cut Latency 40%

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
Enter fullscreen mode Exit fullscreen mode

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])
Enter fullscreen mode Exit fullscreen mode

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=[...],
        )
Enter fullscreen mode Exit fullscreen mode

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]
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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)

Collapse
 
max_quimby profile image
Max Quimby

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_2047 sits 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?