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

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:

1. "What are the hardware and software system requirements for the application?"
2. "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

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.

1. 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).
2. Execution: The vector search queries the database using the small chunks (yielding highly precise matches because small texts contain denser, less diluted semantic vectors).
3. 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.