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

Our production config by document type:

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 } \n Context: { 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):

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

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