Building and Evaluating a Search Engine Over Malware-Classification Papers

A capstone project for AI-Powered Search: BM25 to generative retrieval

llm
BM25
Opensearch
Semantic Ids
vector search
Author

Vinay Pandya

Published

August 4, 2026

Overview

This project builds a search engine over a corpus of arXiv malware-classification papers, implementing and comparing seven retrieval methods head-to-head: BM25 → dense (SPECTER2) → hybrid (RRF) → cross-encoder rerank → late-interaction (ColBERT via Qdrant) → k-means cluster-filtered search → generative retrieval (a trained query→semantic-ID model).

Every method is evaluated on the same 40-query, LLM-judged relevance-labeled eval set, using the same NDCG@10 / MRR metrics, and benchmarked on the same latency harness. The goal was not just to build each technique, but to honestly report where each one wins, where it loses, and why — including findings that didn’t come out the way theory alone would predict.

Code: github.com/vinayhpandya/search-algorithms

Corpus and Setup

  • Domain: arXiv papers on malware classification (cs.CR, cs.LG, cs.CV categories), collected via the arxiv client library across 16 search queries, deduplicated
  • Enrichment: citation counts from Semantic Scholar (used later as an LTR feature)
  • Index: OpenSearch, single-node local Docker setup — chosen because it supports BM25 and k-NN vector search natively in one engine, letting lexical and semantic search share the same index and fuse without cross-system round trips
  • Scope decision: abstract-only (no full-text PDFs), to keep the repo small and the project achievable in a weekend — noted as future work

Methodology: Evaluation

Relevance labeling

Manually labeling enough query-paper pairs for a meaningful eval set (857+ pairs across 40 queries) isn’t realistic by hand in the available time. Instead:

  1. A researcher LLM rates each pair 0/1/2 with a rationale
  2. A critic LLM reviews the researcher’s rating and either confirms or overrides it — final label = critic’s call
  3. A human spot-check on a stratified sample (~65 pairs, weighted toward researcher/critic disagreements) validates the LLM labels

This is standard practice for evaluating relevance judgments at scale (the same logic behind TREC-style IR evaluation and industry search-quality teams checking LLM/crowd-worker judgments) — not a shortcut.

Known calibration issue: only ~3% of labeled pairs were rated “not relevant” (0). Because candidate pools are built from real search results (already topically adjacent to the query), the researcher/critic prompts turned out lenient about what counts as irrelevant. This doesn’t invalidate the labels — NDCG@10 still discriminates well since it weights label=2 three times as heavily as label=1 — but it does mean MRR at the standard threshold (label ≥ 1) saturates near 1.0 for every method and carries no signal. Reporting MRR at threshold ≥ 2 (“how often is the very best result found immediately”) restored real discrimination between methods without needing to relabel anything. (A tighter, binary relevance rubric was designed as a follow-up fix but not run before this writeup — noted under Future Work.)

Metrics

  • NDCG@10: rewards ranking highly-relevant results near the top; uses the full graded (0/1/2) label
  • MRR(≥2): reciprocal rank of the first highly relevant result — the metric that actually discriminates given this corpus’s label distribution
  • Latency: mean / p50 / p95 across 40 queries, local machine, no GPU

Results: Retrieval Quality

Method NDCG@10 MRR(≥1) MRR(≥2)
BM25 0.9003 1.0000 0.5969
Dense 0.9026 1.0000 0.5763
Hybrid 0.9142 1.0000 0.7261
Rerank 0.9131 1.0000 0.7504
Late-interaction 0.9003 1.0000 0.6951
Semantic ID (k-means, cluster-filtered) 0.9624 0.9625 0.2744
Generative (RQ-VAE + trained query→ID model) 0.9005 1.0000 0.2432

Reading the core four methods

Dense underperforms BM25 on MRR(≥2) (0.5763 vs 0.5969) despite near-identical NDCG@10 — SPECTER2 embeddings surface broadly topical neighbors well, but are less precise than exact keyword matching at pinpointing the single best match at rank 1.

Hybrid and rerank both correct for this, each clearly outperforming either input alone on MRR(≥2) — real evidence of fusion and reranking doing genuine work, not just reshuffling noise.

Late-interaction lands between dense and hybrid (MRR≥2 = 0.6951) — consistent with the theoretical expectation that token-level (MaxSim) matching captures finer-grained relevance than single-vector cosine similarity, without needing full cross-attention.

Rerank is the best quality by both metrics — expected, since a cross-encoder sees the full query and document jointly via cross-attention (strictly more expressive than any similarity-metric-based method here), at the cost of only being feasible on a pre-filtered shortlist, never the whole corpus.

Two different semantic-ID-based methods, two different stories

Both remaining rows use “semantic IDs” but are otherwise unrelated methods, worth not conflating:

  • Semantic ID (k-means): a two-level k-means clustering over SPECTER2 embeddings, used to pre-filter/organize the corpus (IVF-style). Highest NDCG@10 of all seven methods (0.9624) but by far the lowest MRR(≥2) (0.2744) — it reliably finds something topically adjacent, but rarely the single best paper. A fixed, non-learned clustering, not a retrieval model in its own right.
  • Generative retrieval: a genuinely trained model (RQ-VAE + query→ID decoder, detailed below) that predicts a semantic ID directly from the query text and looks up matching papers — the actual “generative retrieval” paradigm this project set out to test.

Generative Retrieval: What Was Actually Built

Beyond the six baseline methods above, a from-scratch generative retrieval pipeline was implemented, following the two-stage design used in TIGER-style semantic ID systems:

Stage 1 — RQ-VAE. A small encoder/decoder (MLP, 768→512→256→32) with a 2-level residual vector quantizer (16 codewords/level, EMA codebook updates, dead-code reset) trained on SPECTER2 embeddings to produce a hierarchical 2-digit semantic ID per paper ([coarse, fine]). Trained on the full 194-paper corpus; qualitative inspection showed genuinely topic-coherent clusters (e.g. one cluster grouped five adversarial-malware-attack papers, another grouped five malware-classification-methods papers).

Stage 2 — query-to-ID model. Since no real user queries exist, ~1,940 synthetic (query, semantic_id) pairs were generated via LLM (10 queries per paper, mixing keyword-close and oblique natural-language phrasings). A frozen MiniLM sentence encoder + small 2-layer transformer decoder was trained to autoregressively predict a paper’s semantic ID digit-by-digit from a query, with beam search at inference time. Train/val split was done by paper, not by row, to prevent different phrasings of the same paper’s queries leaking across the split.

Evaluation methodology. The trained model was wired into the same eval harness as the other six methods: for each of the 40 eval queries, beam search produces ranked candidate semantic IDs, each ID’s matching papers are pulled in beam order to build a ranked top-10 list, and that list is scored with the same NDCG@10/MRR machinery and LLM-judged labels used throughout.

Result: a clear, honest negative finding. Training showed a real generalization gap — training exact-match accuracy climbed to 36% while validation exact-match stayed near 2–5%, i.e. the model was memorizing training-set query phrasings rather than learning to generalize across topics. This is directly reflected in the eval numbers: NDCG@10 (0.9005) looks reasonable only because the label set skews lenient (see calibration note above), while MRR(≥2) (0.2432) — the metric that actually measures whether the single best paper is found — is among the weakest of all seven methods. The most likely cause, and the most actionable finding: 194 papers (165 after the train/val split) is simply too little topic diversity for a 256-way discrete prediction task to generalize from — an architecture problem this isn’t; a data-scale problem it is. Scaling to a larger corpus (e.g. the ~100k-paper ML-ArXiv-Papers subset) is the natural next step to test that hypothesis directly.

Latency

Mean / p50 / p95 across 40 queries, local machine, no GPU:

Method Mean (ms) P50 (ms) P95 (ms)
BM25 2.7 2.6 3.3
Late-interaction 12.6 12.3 16.0
Generative 20.4 20.0 23.9
Semantic ID (k-means) 30.5 30.4 36.0
Dense 42.7 42.7 48.0
Hybrid 59.6 59.4 64.8
Rerank 187.8 186.8 200.9

BM25 is fastest by a wide margin, as expected — no model inference at all. Generative retrieval lands in the middle of the pack: slower than BM25 and late-interaction, but notably faster than dense search despite also requiring a forward pass, since the query encoder here (frozen MiniLM) is far smaller than SPECTER2, and lookup after ID prediction is a cheap dictionary access rather than a k-NN search over the full corpus. Rerank is the clear latency outlier, consistent with cross-encoders’ cost scaling with the number of candidates scored, not just query size — the tradeoff that pairs it with a cheap first-stage retriever in practice rather than running it corpus-wide.

Future Work

  • Rerun evaluation under a binary (0/1), query-specific relevance rubric — designed to fix the labeling calibration issue described above, but not executed before this writeup.
  • Scale the corpus (e.g. to the ~100k-paper ML-ArXiv-Papers subset) to directly test whether generative retrieval’s weak MRR(≥2) is a data-scale problem, as the evidence here suggests, rather than an architectural one.
  • Hypencoders Testing a small Qnet and hype-encoder based network, want to see how fast it is to train and how well the NDCG scores are
  • Full-text PDFs and a RAG QA layer on top of the retrieval stack — scoped out from the very beginning to keep this project achievable in the available time.