📦 Open Source Repository: The complete codebase, architecture schemas, and implementation files for this system are available on GitHub at rishav-dahal/Modular-Query-Refinement.

In the current artificial intelligence boom, the standard response to almost any natural language processing problem is: “Just throw a fine-tuned 7B LLM or a dense vector embedding database at it.”

While dense vector retrieval (RAG) is fantastic, it comes with brutal production trade-offs:

  • Running continuous vector embeddings for millions of queries requires expensive GPU infrastructure.
  • Vector search latency rarely drops below 150ms.
  • LLMs suffer from “semantic drift” or hallucinations, appending tangential keywords that dilute original search intent.

When I was conducting research on information retrieval for our paper presented at the International Conference on Innovation in Computing, Science, Engineering and Technology (ICICSET 2025), we explored a much faster, mathematically grounded alternative: Unsupervised Query Refinement using Latent Dirichlet Allocation (LDA).

Here is why short search queries break traditional search engines, how LDA models latent semantic concepts in documents, and how we built a sub-15ms query refinement pipeline running purely on CPU.

1. The 3-Word Dilemma & Query Drift

In information retrieval, the average user search query is fewer than three words long (e.g., "apple health").

This triggers the classic Vocabulary Mismatch Problem: the exact words the user typed do not appear in the most relevant documents, even though the underlying meaning matches.

Historically, search engines attempted to fix this using Pseudo-Relevance Feedback (PRF):

  1. Run an initial search for "apple health".
  2. Grab the top 10 returned documents.
  3. Extract the most frequent words from those 10 documents and append them to the query.

The Disaster of Naive PRF:

Initial Query: "apple health"
       │
       ▼
[Top-10 Documents Retrieved]
 ├── Doc 1: Apple Watch optical heart rate sensors...
 ├── Doc 2: Health benefits of organic apple orchards...
 ├── Doc 3: Eating fiber to prevent cardiovascular disease...
       ▼
Naive Frequency Extraction Blends Unrelated Concepts:
Expanded Query: "apple health orchard nutrition watch heart" 
Result: Absolute nonsense results (Query Drift)!

Because naive PRF ignores topic boundaries, polysemous words (words with multiple distinct meanings like “Apple” the company vs “apple” the fruit) contaminate the expanded query.

2. Mathematical Intuition: Latent Dirichlet Allocation

Latent Dirichlet Allocation (LDA), formulated by David Blei, Andrew Ng, and Michael Jordan, solves this by treating documents as probabilistic mixtures of latent topics, where each topic is a probability distribution over words.

                  Document Collection
                          │
            Dirichlet Prior Distribution (α)
                          │
                          ▼
            [Topic 1: Consumer Tech] ──► {watch: 0.08, sensor: 0.05, screen: 0.04...}
            [Topic 2: Agriculture]   ──► {orchard: 0.09, fruit: 0.07, harvest: 0.04...}
            [Topic 3: Cardiology]    ──► {heart: 0.09, blood: 0.06, pressure: 0.05...}

Instead of blindly taking word frequencies across all returned documents, our system:

  1. Projects the initial query onto the trained LDA topic space.
  2. Identifies the single most salient topic distribution $\theta_q = P(Topic | Query)$.
  3. Restricts query expansion words strictly to terms that have high probability $P(Term | Topic)$ within that dominant topic.

If the user is asking about Apple Watch sensors, our algorithm expands with tech hardware terms; if they’re asking about fruit fiber, it expands with dietary terms—never both simultaneously.

3. Production Python Implementation with Scikit-Learn

Here is the clean, CPU-friendly implementation using sklearn.decomposition.LatentDirichletAllocation:

# lda_refiner.py
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.decomposition import LatentDirichletAllocation

class LDAQueryRefiner:
    def __init__(self, n_topics=10, max_features=5000):
        self.n_topics = n_topics
        self.vectorizer = TfidfVectorizer(
            max_df=0.90,
            min_df=2,
            stop_words="english",
            max_features=max_features
        )
        # Online variational Bayes for fast training and low memory footprint
        self.lda = LatentDirichletAllocation(
            n_components=n_topics,
            learning_method="online",
            random_state=42,
            batch_size=128
        )
        self.feature_names = None

    def fit(self, corpus_documents):
        """Fits TF-IDF and LDA on document collection"""
        tfidf = self.vectorizer.fit_transform(corpus_documents)
        self.lda.fit(tfidf)
        self.feature_names = np.array(self.vectorizer.get_feature_names_out())
        print(f"[FIT] Trained {self.n_topics} latent topics across {len(corpus_documents)} documents.")

    def refine_query(self, raw_query: str, top_n_terms: int = 3) -> str:
        """
        Projects short query into topic space and appends 
        highest-probability coherent terms from dominant topic.
        """
        query_vec = self.vectorizer.transform([raw_query])
        
        # Check if query terms exist in vocabulary
        if query_vec.nnz == 0:
            return raw_query # Return original if out of vocabulary

        # Predict topic distribution for query: shape (1, n_topics)
        topic_distribution = self.lda.transform(query_vec)[0]
        dominant_topic_idx = np.argmax(topic_distribution)

        # Get term weights for this dominant topic
        topic_term_weights = self.lda.components_[dominant_topic_idx]
        
        # Exclude terms that were already in the user's query
        query_words = set(raw_query.lower().split())
        sorted_indices = np.argsort(topic_term_weights)[::-1]

        expansion_terms = []
        for idx in sorted_indices:
            term = self.feature_names[idx]
            if term not in query_words:
                expansion_terms.append(term)
                if len(expansion_terms) >= top_n_terms:
                    break

        refined_query = f"{raw_query} {' '.join(expansion_terms)}"
        return refined_query

# Example verification:
if __name__ == "__main__":
    sample_docs = [
        "Apple watch series tracks heart rate ECG and physical fitness vitals.",
        "Cardiovascular exercise reduces resting heart rate and blood pressure.",
        "Organic apple orchards produce crisp fruit harvested in autumn for nutrition.",
        "Smartphones and wearable smartwatches feature OLED displays and Bluetooth.",
        "Dietary fiber from fruits like apples and oranges aids gut digestion."
    ]

    engine = LDAQueryRefiner(n_topics=3)
    engine.fit(sample_docs)

    query = "apple watch"
    refined = engine.refine_query(query, top_n_terms=2)
    print(f"Original: '{query}' ──► Refined: '{refined}'")

4. Benchmark Results & Key Takeaways

In our experimental evaluation against benchmark information retrieval corpora:

  1. Mean Average Precision (MAP): LDA-based topic refinement improved MAP by 14.2% over unexpanded baseline BM25 queries.
  2. Zero Semantic Drift: Because expansion words are constrained to the single highest-probability Dirichlet distribution, cross-domain topic pollution dropped by 82% compared to standard Rocchio feedback.
  3. Execution Speed: Inference takes under 12 milliseconds on a standard single-core CPU, making it suitable for high-traffic search bars where spinning up a GPU cluster is financially unjustifiable.

Summary

Before committing to heavy LLM pipelines with recurring API bills and 200ms latency penalties, consider classic probabilistic topic modeling. Latent Dirichlet Allocation provides deterministic, mathematically sound query enhancement that is lightweight, interpretable, and blazing fast on production systems.

🛠️ GitHub Repository & Next Steps

The complete open-source source code and architecture discussed in this guide are publicly available:

If you’re building a similar system or encounter edge cases in your deployment, feel free to star the repo, file an issue, or submit an optimization pull request!