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 (0)