Skip to content
Ashish.
All posts
Diagram showing query rewriting process in a RAG pipeline.
6 min readDevelopmentDevelopers, ML EngineersFeatured#rag#query-rewriting#nlp#search#ai#machine-learning

Query Rewriting: The Underrated RAG Upgrade

Explore how query rewriting and expansion improve RAG recall by resolving coreference and enhancing conversational search capabilities.

By Ashish KumarPart 9 of RAG from Zero to Production

In a standard Retrieval-Augmented Generation (RAG) pipeline, the flow is typically linear: a user asks a question, the system retrieves relevant documents, and a Large Language Model (LLM) generates an answer. While this simplicity is appealing, it often breaks down when users engage in natural, multi-turn conversations. In these scenarios, queries rarely stand alone; they rely on implicit context, layered reasoning, and referential ambiguity that naive vector search cannot parse.

The missing link in many production-grade systems is query rewriting. This process transforms raw user input into a form optimized for retrieval before it hits the vector database. Without it, systems suffer from low recall because the embedding of a vague follow-up like "What was the budget?" differs significantly from "What was the 2023 marketing budget?" even if the answer resides in the same document.

The Coreference Trap

Consider a typical conversation thread within a business intelligence context:

  1. User: "Summarize the Q3 financial report."
  2. Assistant: "The report shows a 15% increase in revenue..."
  3. User: "What about the expenses?"

A naive RAG system takes the literal string "What about the expenses?" and embeds it directly. The vector space likely returns documents about general expense policies or unrelated cost centers, missing the specific Q3 financial report entirely. The system fails because the query lacks the necessary entity anchors.

Query rewriting resolves this via coreference resolution. The system analyzes the conversation history to identify that "the expenses" refers to "Q3 financial report expenses." It rewrites the query to: "What were the expenses in the Q3 financial report?"

This transformation shifts the query vector closer to the embedding of the target document. The mechanism here is contextual grounding. By explicitly binding pronouns to their antecedents, we reduce the semantic distance between the user's intent and the stored knowledge.

Expanding Semantic Breadth

Coreference resolution handles reference, but query expansion handles ambiguity. Users often use shorthand or domain-specific jargon that may not match the terminology used in the source documents.

For example, a developer might ask: "How do I fix a segfault in Python?"

The source documents might use the term "segmentation fault" or "memory access violation." A direct vector match might fail if the embedding model treats "segfault" and "segmentation fault" as distinct concepts due to lack of co-occurrence in the training data.

Query expansion addresses this by generating a richer query. An LLM can be prompted to add synonyms, related concepts, and broader categories. The rewritten query might become: "How to resolve segmentation faults or memory access violations in Python code."

This expansion increases the likelihood of recall (finding the relevant document) without necessarily sacrificing precision (finding irrelevant noise), provided the expansion is constrained. Over-expansion can introduce noise, so techniques like HyDE (Hypothetical Document Embeddings) are often employed. In HyDE, the system generates a hypothetical answer to the query, embeds that answer, and uses its vector to retrieve similar documents. This bridges the gap between the sparse, informal nature of user queries and the dense, formal nature of technical documentation.

Decomposing Multi-Hop Queries

The most significant limitation of simple retrieval is its inability to handle multi-hop questions. These are queries that require synthesizing information from multiple distinct sources.

Example: "Who is the CEO of the company that acquired Instagram in 2012?"

A naive retrieval system might search for "CEO of company acquired Instagram" and fail, because no single document contains both "CEO" and "acquired Instagram" in the relevant context. The system needs to break the problem down.

Query rewriting here acts as a decomposition engine. The system rewrites the single complex query into a sequence of sub-queries:

  1. "Which company acquired Instagram in 2012?"
  2. "Who is the CEO of [Company Name]?"

This approach, often seen in Self-Ask or Step-back prompting frameworks, allows the retrieval system to fetch intermediate facts. The first rewrite retrieves the company name. The second rewrite uses that fact to find the CEO.

This mechanism transforms a retrieval failure into a multi-step reasoning process. It requires the rewriting module to maintain state across steps, injecting the results of previous retrievals into subsequent query rewrites.

Implementation Trade-offs

Implementing query rewriting adds latency and cost. Each rewrite requires an LLM call, which is slower and more expensive than a simple vector lookup. Therefore, not all queries need rewriting.

A practical architecture uses a classifier to decide when to rewrite. Simple factual queries ("What is the capital of France?") can go straight to retrieval. Complex, ambiguous, or multi-turn queries trigger the rewriting module. This hybrid approach balances performance with cost.

Furthermore, the quality of the rewritten query depends heavily on the prompt engineering of the rewriting LLM. Poor prompts can lead to hallucinated constraints or irrelevant expansions.

Common Pitfalls

  1. Over-expansion leading to noise: Adding too many synonyms or broadening the query too much can dilute the signal, causing the retriever to pull in irrelevant documents that share keywords but not intent.
  2. Latency/cost implications of LLM calls: Every rewrite step adds round-trip time. Unnecessary rewrites for simple queries can degrade user experience and increase infrastructure costs significantly.
  3. Hallucination risks from poor prompt engineering: If the rewriting model is not strictly constrained, it may invent entities or relationships that do not exist in the context, leading to retrieval of completely fabricated or unrelated content.

Practical Takeaways

  1. Use a classifier to gate rewriting: Do not rewrite every query. Route simple, standalone questions directly to the vector store and reserve rewriting for ambiguous, multi-turn, or complex queries.
  2. HyDE bridges semantic gaps: When standard keyword or vector matching fails due to terminology mismatches, use Hypothetical Document Embeddings to generate a synthetic document that aligns better with your index.
  3. Decompose multi-hop queries: Break complex questions into sub-queries to allow the system to retrieve intermediate facts, ensuring that no single retrieval step is overwhelmed by complexity.

FAQ

What is coreference resolution in RAG? Coreference resolution is the process of identifying when different words or phrases refer to the same entity. In RAG, it allows the system to replace pronouns (like "it," "they," or "the budget") with their full antecedents from previous conversation turns, enabling accurate retrieval in multi-turn dialogues.

How does HyDE work? HyDE (Hypothetical Document Embeddings) works by having an LLM generate a hypothetical answer or document based on the user's query. The system then embeds this hypothetical text and uses its vector to search the database. This often yields better results than embedding the raw query because the hypothetical text mimics the style and density of the indexed documents.

When should query rewriting be used? Query rewriting should be used when the user's input is ambiguous, relies on previous context (multi-turn), or contains complex logical requirements (multi-hop). It is less necessary for simple, standalone factual questions where direct vector matching is sufficient.

Related posts