DEV Community

Imus
Imus

Posted on

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

#llm #ai #evaluation #agents

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

Collapse
 
mads_hansen_27b33ebfee4c9 profile image
Mads Hansen •

Bayesian search is a good fit here, but it creates a subtle evaluation trap: if the same 200-query golden set selects hyperparameters and reports the final 95% recall, the optimizer can overfit that set just like a model overfits training data.

I would split queries into optimization, validation, and untouched test sets, grouped by source document or customer so near-duplicate questions cannot leak across splits. Report bootstrap confidence intervals and per-slice recall for exact identifiers, semantic paraphrases, multi-hop questions, and no-answer cases; the aggregate can hide a serious regression in one class. Latency should also be measured under fixed concurrency with warm/cold cache states and tail percentiles. One more instrumentation nuance: production recall cannot be inferred from the retriever’s own result without independent relevance labels, so sampled traffic needs human judgments, clicks with caveats, or planted canary queries. Otherwise the dashboard may confidently measure agreement with itself.

Collapse
 
topstar_ai profile image
Luis Cruz •

The use of a hybrid retrieval approach, combining BM25, vector search, and reranking, is particularly intriguing, as it addresses the limitations of relying solely on vector search or BM25. I've found that in my own projects, incorporating a similar hybrid approach can significantly improve the recall and precision of search results. The addition of a cross-encoder rerank step, which boosts correlation with relevance to approximately 0.92, is also a valuable insight - at what point do you find the cost of this additional step outweighs the benefits in terms of improved recall, especially considering the 50ms cost for the 50→5 funnel?