Advanced RAG: Better Retrieval with Re-ranking and Query Expansion | FromZero2AI Blog

A simple "Naive" Retrieval-Augmented Generation (RAG) system consists of chunking documents, embedding them into a vector database, querying the database, and passing the top K results directly to a Large Language Model (LLM).

While Naive RAG works well for basic demos, it frequently fails in production. Standard user queries are often short, poorly phrased, or lacking the technical vocabulary contained in your enterprise documentation. Furthermore, vector databases occasionally return irrelevant chunks that dilute the LLM's context window, causing the model to miss critical details—a phenomenon known as the "Lost in the Middle" problem.

To solve this, advanced architectures modify both the input query and the retrieved results before feeding them to the generation stage. This article explores two major components of Advanced RAG: Query Expansion and Cross-Encoder Re-ranking.

1. Query Expansion: Expanding the Search Net

Query expansion is the process of generating multiple variations of a user's query to maximize the probability of matching relevant documents. There are two primary techniques:

I. Sub-Query / Multi-Query Retrieval

User queries can be ambiguous or multi-faceted. If a user asks: "What are the system requirements for our application and how do we deploy it to AWS?", the vector database search might struggle to find chunks covering both topics in a single retrieval step.

Using an LLM, we can decompose the query into distinct sub-questions:

"What are the hardware and software system requirements for the application?"
"What are the deployment steps to host the application on AWS?"

We execute vector searches for each generated query, combine the retrieved documents, and remove duplicates.

II. Query Rewriting (HyDE - Hypothetical Document Embeddings)

HyDE takes the user's query and asks an LLM to generate a hypothetical answer. Even if the hypothetical answer contains hallucinated facts, its vocabulary, phrasing, and writing style will be much closer to the source documents in the vector database than the raw user query.

Advanced RAG Retrieval: Advanced RAG architecture diagram showing query expansion, vector DB query, cross-encoder re-ranking, and context generation Advanced RAG Retrieval: Advanced RAG architecture diagram showing query expansion, vector DB query, cross-encoder re-ranking, and context generation

By embedding and searching with this hypothetical answer, we retrieve actual documents that match the semantic structure of a resolved answer.

2. Re-ranking: Mitigating the "Lost in the Middle" Problem

Researchers have shown that LLMs are excellent at extracting information from the very beginning or the very end of their input context, but they struggle to process facts hidden in the middle of a massive block of text.

If you retrieve 20 documents to ensure high recall, but only 3 contain the answer, presenting all 20 to the LLM increases cost, latency, and the likelihood of hallucination.

Re-ranking inserts a gatekeeper model between the database and the LLM.

You query the database for K=50 candidates.
You pass these 50 candidates through a specialized Cross-Encoder model, which scores the literal alignment between the query and each chunk.
You select only the top N=4 highest-scoring chunks to send to the generator LLM.

This keeps the LLM's prompt concise, high-density, and free of noisy, irrelevant information.

3. Advanced RAG Implementation in Python

The following script implements a complete retrieval pipeline using Multi-Query expansion (simulating LLM query rewriting) and Cross-Encoder re-ranking.

import os
from typing import List
from sentence_transformers import SentenceTransformer, CrossEncoder, util
import numpy as np

# Initialize embedding and re-ranking models
embed_model = SentenceTransformer("all-MiniLM-L6-v2")
rerank_model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

# Knowledge base chunks
knowledge_base = [
    "AWS deployments require Docker containers configured with ECS or EKS services.",
    "The system requires at least 16GB of RAM and a quad-core CPU to run the analytics engine.",
    "For local development, install Docker Desktop and run the startup shell script.",
    "Python 3.10 or higher is required to run the pipeline libraries.",
    "To deploy on AWS, set up an IAM user with EC2 Container Registry write access.",
    "Our backup policy requires daily snapshots of all RDS databases with a 30-day retention."
]

# Create local index
document_embeddings = embed_model.encode(knowledge_base, convert_to_tensor=True)

def simulate_llm_query_expansion(original_query: str) -> List[str]:
    """
    Simulates an LLM call generating alternative formulations of a query.
    In production, this would be an actual LLM API call using a prompt template.
    """
    # Let's say the user asked about running the app on AWS.
    # The LLM expands this into technical keywords and deployment terms.
    if "AWS" in original_query or "deploy" in original_query:
        return [
            original_query,
            "AWS deployment configurations and container services ECS EKS",
            "cloud infrastructure setup for application deployment",
            "how to host the system on Amazon Web Services"
        ]
    return [original_query]

def advanced_retrieve(query: str, top_k_initial: int = 4, top_n_final: int = 2) -> List[str]:
    print(f"Original User Query: '{query}'")
    
    # Step 1: Query Expansion
    expanded_queries = simulate_llm_query_expansion(query)
    print(f"Expanded Queries: {expanded_queries}\n")
    
    # Step 2: Parallel Retrieval and De-duplication
    retrieved_indices = set()
    for q in expanded_queries:
        q_emb = embed_model.encode(q, convert_to_tensor=True)
        scores = util.cos_sim(q_emb, document_embeddings)[0]
        # Get top matching indices for this sub-query
        top_idx = np.argsort(scores.cpu().numpy())[::-1][:top_k_initial]
        for idx in top_idx:
            retrieved_indices.add(int(idx))
            
    candidate_chunks = [knowledge_base[i] for i in retrieved_indices]
    print(f"Candidates Retrieved (Count: {len(candidate_chunks)}):")
    for doc in candidate_chunks:
        print(f" - {doc}")
        
    # Step 3: Cross-Encoder Re-ranking
    pairs = [[query, doc] for doc in candidate_chunks]
    rerank_scores = rerank_model.predict(pairs)
    
    # Sort candidates by re-rank score
    ranked_indices = np.argsort(rerank_scores)[::-1]
    
    print("\n--- Final Re-ranked Context sent to Generator LLM ---")
    final_context = []
    for rank in ranked_indices[:top_n_final]:
        score = rerank_scores[rank]
        matched_doc = candidate_chunks[rank]
        print(f"Score: {score:.4f} | Chunk: {matched_doc}")
        final_context.append(matched_doc)
        
    return final_context

if __name__ == "__main__":
    context = advanced_retrieve("How do we deploy this application on AWS?", top_k_initial=2, top_n_final=2)

4. Chunking Optimizations: Parent-Child Retrieval

To further optimize advanced RAG, consider splitting retrieval from generation chunks.

Small-to-Large Chunking: In this model, you chunk your text into small sentences or segments (e.g., 50 tokens) and index them. Each small chunk is linked to a larger parent chunk (e.g., 500 tokens).
Execution: The vector search queries the database using the small chunks (yielding highly precise matches because small texts contain denser, less diluted semantic vectors).
Re-routing: Once a small chunk is retrieved, the pipeline fetches its parent chunk from the database and inserts that parent chunk into the LLM prompt. This gives the generator LLM the full contextual background needed to answer, without suffering from sparse keyword matches.

Where This Gets Expensive

Every technique above costs something, and the bill is worth stating plainly before you adopt all of them at once.

Query expansion multiplies your retrieval work. Four generated queries means four embedding calls and four vector searches per user request, plus an LLM call to produce the variants — and that LLM call sits on the critical path before any retrieval starts. On a system where users expect a first token in under a second, that alone can blow the budget. HyDE is worse: it needs a full generation, not a short rewrite.

Re-ranking is cheaper than it looks on paper but slower than people expect in practice. A cross-encoder scores every query-document pair independently, so 50 candidates means 50 forward passes with no batching shortcut across queries. On CPU that is often 200-500ms. Hosted re-rankers move the latency to a network hop, which may or may not be an improvement.

The honest test is whether you actually have a retrieval problem. If you evaluate your pipeline and find the correct chunk is already in the top 5 the vast majority of the time, re-ranking has almost nothing to fix, and your failures are coming from the generation step or from gaps in the source documents — neither of which these techniques touch. Measure retrieval recall on a labeled query set first. Teams routinely add both query expansion and re-ranking, triple their per-query cost, and discover the remaining errors were chunking artifacts all along: a table split across two chunks, or a policy document where the exception clause landed in a different segment than the rule it modifies.

If you adopt one thing from this article, make it the evaluation set, not the re-ranker.