Advanced RAG: Query Expansion and Multi-Query Retrieval
Retrieval-Augmented Generation (RAG) has become the backbone of modern LLM applications, but vanilla RAG pipelines often stumble when users submit short, ambiguous, or poorly phrased queries. A single embedding computed from a terse user prompt may not surface the most relevant documents, simply because the user's wording doesn't match the language used in the source corpus. Query expansion and multi-query retrieval are two complementary techniques that address this gap by enriching and broadening the retrieval step before generation happens.
What Is Query Expansion?
Query expansion is the process of augmenting an original user query with additional terms, synonyms, paraphrases, or related concepts before performing retrieval. Instead of embedding just the raw user input, the system embeds an enriched version of the query — or several enriched versions — to increase the chance of retrieving semantically relevant chunks.
There are several flavors of query expansion:
- Lexical expansion: Adding synonyms or domain terms using a thesaurus or controlled vocabulary.
- LLM-based expansion: Asking an LLM to rewrite or expand the query with additional context.
- Pseudo-relevance feedback: Retrieving an initial set of documents, extracting terms from them, and re-running retrieval with the expanded query.
- Hypothetical document embeddings (HyDE): Generating a hypothetical answer and using it as the query.
What Is Multi-Query Retrieval?
Multi-query retrieval takes query expansion one step further. Instead of producing a single expanded query, the system generates multiple distinct reformulations of the original query and runs retrieval for each one. The resulting document sets are then merged, deduplicated, and ranked before being passed to the generator. This dramatically improves recall because different phrasings can surface different relevant chunks.
The typical multi-query pipeline looks like this:
- Receive the original user query.
- Use an LLM to generate N alternative phrasings of the query.
- Embed and retrieve documents for each of the N queries.
- Union the retrieved documents and deduplicate them.
- Optionally re-rank the combined set using a cross-encoder or another scoring method.
- Pass the top-K documents to the generator LLM.
Why It Matters
Standard RAG assumes that the user's query is a good proxy for the documents they want to retrieve. In practice, this assumption breaks down frequently. A user might ask "How do I handle auth failures?" while the documentation uses the phrase "authentication error handling." A single embedding comparison may miss the connection. Multi-query retrieval and query expansion mitigate this vocabulary mismatch problem and improve robustness against:
- Ambiguous or underspecified queries.
- Domain-specific jargon that users may not know.
- Multi-part questions where different sub-questions need different documents.
- Cross-lingual or mixed-language queries.
The trade-off is latency and cost: you perform more embedding calls and more vector searches per user turn. For most production systems, the retrieval quality gains are well worth it.
How to Use It: A Practical Implementation
Below is a complete, runnable example using Python, OpenAI's API, and a simple in-memory vector store. The same pattern applies whether you use LangChain, LlamaIndex, or a custom pipeline.
import os
import openai
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
openai.api_key = os.environ["OPENAI_API_KEY"]
# --- A tiny in-memory vector store for demonstration ---
class SimpleVectorStore:
def __init__(self):
self.documents = []
self.embeddings = []
def add(self, docs):
for doc in docs:
emb = get_embedding(doc)
self.documents.append(doc)
self.embeddings.append(emb)
def search(self, query, k=3):
q_emb = get_embedding(query)
scores = cosine_similarity([q_emb], self.embeddings)[0]
top_idx = np.argsort(scores)[::-1][:k]
return [(self.documents[i], scores[i]) for i in top_idx]
def get_embedding(text, model="text-embedding-3-small"):
resp = openai.embeddings.create(input=text, model=model)
return resp.data[0].embedding
# --- Sample corpus ---
corpus = [
"Authentication errors usually occur when the token has expired or is invalid.",
"To reset your password, navigate to Settings > Security > Reset Password.",
"OAuth 2.0 supports authorization code, implicit, and client credentials flows.",
"Rate limiting is enforced at 100 requests per minute per API key.",
"If login fails repeatedly, check that your account is not locked by an administrator.",
"JWT tokens should be signed with a strong secret and rotated periodically.",
]
store = SimpleVectorStore()
store.add(corpus)
Now let's implement the multi-query generation step. We ask the LLM to produce several alternative phrasings of the user's question.
def generate_multi_queries(user_query, n=3):
prompt = (
f"You are an expert search assistant. "
f"Rewrite the following user query into {n} different phrasings "
f"that a user might use to search for the same information. "
f"Return them as a numbered list, one per line, with no extra commentary.\n\n"
f"Original query: {user_query}\n\n"
f"Alternative queries:"
)
resp = openai.chat.completions.create(
model="gpt-4o-mini",
messages=[{"role": "user", "content": prompt}],
temperature=0.7,
)
text = resp.choices[0].message.content.strip()
queries = []
for line in text.splitlines():
line = line.strip()
if line and line[0].isdigit():
# Strip the "1. " prefix
cleaned = line.split(".", 1)[1].strip()
queries.append(cleaned)
return queries
With the multi-query generator in place, we can build the full retrieval pipeline. We retrieve documents for each query, merge the results, deduplicate, and keep the top-K by maximum score.
def multi_query_retrieve(user_query, store, n_queries=3, k_per_query=3, final_k=4):
# Step 1: Generate alternative queries
alt_queries = generate_multi_queries(user_query, n=n_queries)
all_queries = [user_query] + alt_queries
print("Generated queries:")
for q in all_queries:
print(f" - {q}")
# Step 2: Retrieve for each query
seen = {}
for q in all_queries:
results = store.search(q, k=k_per_query)
for doc, score in results:
if doc not in seen or score > seen[doc]:
seen[doc] = score
# Step 3: Rank by best score and take top-K
ranked = sorted(seen.items(), key=lambda x: x[1], reverse=True)[:final_k]
return [doc for doc, _ in ranked]
# --- Run it ---
user_query = "How do I handle auth failures?"
retrieved_docs = multi_query_retrieve(user_query, store)
print("\nFinal retrieved documents:")
for i, doc in enumerate(retrieved_docs, 1):
print(f"{i}. {doc}")
The output will look something like this:
Generated queries:
- How do I handle auth failures?
- What should I do when authentication errors occur?
- Why does login keep failing and how can I fix it?
- How to troubleshoot invalid or expired authentication tokens?
Final retrieved documents:
1. Authentication errors usually occur when the token has expired or is invalid.
2. If login fails repeatedly, check that your account is not locked by an administrator.
3. JWT tokens should be signed with a strong secret and rotated periodically.
4. To reset your password, navigate to Settings > Security > Reset Password.
Notice how the alternative phrasings surfaced documents that a single embedding of "How do I handle auth failures?" might have ranked lower. The merged result set covers token issues, account lockouts, and password resets — all relevant facets of the original question.
Adding HyDE as a Query Expansion Variant
Hypothetical Document Embeddings (HyDE) is a powerful query expansion technique where the LLM generates a plausible answer to the query, and that hypothetical answer is used as the search vector. The idea is that a generated answer, even if factually imperfect, will use vocabulary closer to the corpus than the terse question.
def hyde_query(user_query, n_hypotheticals=1):
prompt = (
f"Please write a short paragraph (3-5 sentences) that would be a good answer "
f"to the following question. Do not worry about factual accuracy; focus on "
f"using natural, descriptive language.\n\nQuestion: {user_query}\n\nAnswer:"
)
hypotheticals = []
for _ in range(n_hypotheticals):
resp = openai.chat.completions.create(
model="gpt-4o-mini",
messages=[{"role": "user", "content": prompt}],
temperature=0.8,
)
hypotheticals.append(resp.choices[0].message.content.strip())
return hypotheticals
def hyde_retrieve(user_query, store, final_k=4):
hyp_docs = hyde_query(user_query, n_hypotheticals=2)
seen = {}
for hyp in hyp_docs:
for doc, score in store.search(hyp, k=3):
if doc not in seen or score > seen[doc]:
seen[doc] = score
ranked = sorted(seen.items(), key=lambda x: x[1], reverse=True)[:final_k]
return [doc for doc, _ in ranked]
You can even combine HyDE with multi-query retrieval: generate alternative queries, generate hypothetical answers for each, and retrieve against all of them. This hybrid approach tends to produce the highest recall on complex question-answering benchmarks.
Adding a Re-Ranking Step
After merging documents from multiple queries, a re-ranking step using a cross-encoder can significantly improve precision. Cross-encoders score query-document pairs jointly, capturing interactions that bi-encoder embeddings miss.
from sentence_transformers import CrossEncoder
# Load a lightweight cross-encoder
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
def rerank(user_query, documents, top_k=4):
pairs = [(user_query, doc) for doc in documents]
scores = reranker.predict(pairs)
ranked = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)
return [doc for doc, _ in ranked[:top_k]]
# Full pipeline
candidates = multi_query_retrieve(user_query, store, final_k=10)
final_docs = rerank(user_query, candidates, top_k=4)
Best Practices
- Start simple. Begin with 3-5 alternative queries. More queries yield diminishing returns and increase latency.
- Cache aggressively. Cache both the generated alternative queries and their embeddings. Many user queries repeat, especially in support and FAQ scenarios.
- Control temperature. Use a moderate temperature (0.5-0.8) for query generation so alternatives are diverse but not nonsensical.
- Deduplicate carefully. Use exact string matching or a similarity threshold to avoid feeding the generator near-duplicate chunks.
- Always include the original query. Never discard the user's exact phrasing — it sometimes matches best.
- Re-rank when possible. Multi-query retrieval increases recall; re-ranking restores precision. The two techniques are complementary.
- Monitor cost and latency. Each additional query means one more LLM call, one more embedding call, and one more vector search. Measure the end-to-end impact in production.
- Log generated queries. Inspecting the alternative queries your LLM produces is invaluable for debugging retrieval failures and tuning prompts.
- Consider query classification. Not every query needs expansion. Simple, unambiguous lookups can skip the multi-query step to save latency.
Conclusion
Query expansion and multi-query retrieval are among the highest-leverage improvements you can make to a RAG pipeline. By acknowledging that a single user phrasing is rarely sufficient to capture all relevant documents, these techniques trade a modest amount of latency and cost for substantial gains in recall and answer quality. When combined with deduplication, cross-encoder re-ranking, and careful prompt engineering, multi-query retrieval transforms a fragile single-shot retriever into a robust system that gracefully handles the messy, ambiguous queries real users actually submit. Start with a small number of alternative queries, measure retrieval quality on your own evaluation set, and iterate from there.