Optimizing MIPS Retrieval with Learned Networks
Markdown--- name: optimizing-mips-retrieval-with-learned-networks description: Designs and trains small neural networks (SupportNet or KeyNet) to accelerate Maximum Inner Product Search (MIPS) by amortizing compute offline on a known query distribution, rather than tuning generic query-agnostic index parameters. Use when a RAG/retrieval system has a large mostly-static key database, a stable/known query distribution, and a latency-vs-recall bottleneck in its vector index (e.g., IVF cluster routing or pre-index query shaping). ---
Given: an IVF index with k clusters, a production query log, and a latency budget.
- Sample N queries from the real deployment query distribution (not the key distribution).
- Run exhaustive MIPS once per query to get ground truth: either per-cluster max inner product (for SupportNet) or the best-matching key vector (for KeyNet).
- Train a small NN on this (query → label) dataset.
- At inference, use the NN's output to route to top-s clusters (SupportNet) or to rewrite the query before feeding it into the existing index (KeyNet).
- Benchmark recall-vs-latency and recall-vs-budget against centroid/raw-query baselines at matched budgets.
- Pick the smallest model size (XS/S/M/L as % of database size) that hits the target recall within the latency budget.
Core reframe: MIPS search = learning the support function of the key set (a convex function whose value is max inner product, whose gradient is the argmax key). Learn this function (or its argmax) for the actual query distribution instead of tuning generic geometric index knobs.
Progress:
- Step 1: Characterize the setup — key database, query distribution, existing index type
- Step 2: Choose intervention point — routing (SupportNet) vs. query rewriting (KeyNet)
- Step 3: Build supervised training set via exhaustive ground-truth search
- Step 4: Train the model
- Step 5: Benchmark recall-vs-latency and recall-vs-budget vs. baseline
- Step 6: Select model size for the latency budget
Step 1 — Characterize the setup. Check whether keys and queries come from different distributions (e.g., short questions vs. long documents). This mismatch is exactly where query-agnostic geometric methods (pure clustering/quantization) underperform and where learned methods have the most upside. Note the existing index type (IVF/clustered vs. flat) and whether you can modify its internals.
Step 2 — Choose intervention point.
- Routing (SupportNet): replaces "compare query to k centroids" in an IVF-style index. Use when you control index internals.
- Query rewriting (KeyNet): sits upstream of an existing, already-tuned/black-box index. Use when you can't or don't want to modify the index.
Step 3 — Build training data. For each sampled query (from the real/expected query distribution):
- Run exhaustive ground-truth MIPS once (expensive, but offline/amortized, done once per training query).
- Store (query, optimal per-cluster score) for SupportNet, or (query, optimal key vector) for KeyNet.
Step 4 — Train the model.
- SupportNet: loosely-constrained input-convex neural network (ICNN) to exploit known convexity of the support function. Output k scores jointly (one per cluster), trained via MSE regression against ground-truth support values.
- KeyNet: standard regression network, query → best-matching key vector. No hard convexity constraint; add a consistency loss derived from Euler's theorem for homogeneous functions to keep it geometrically grounded.
Step 5 — Benchmark on two orthogonal axes:
- Recall vs. wall-clock latency (include NN inference cost — never ignore model overhead).
- Recall vs. search budget (fraction of database/clusters scanned) — isolates algorithmic routing quality from hardware/implementation effects.
- Always compare against centroid routing / raw-query IVF search at matched budgets.
Step 6 — Pick model size. Sweep sizes as % of database size (e.g., 1/5/10/20% → XS/S/M/L). Choose the smallest model hitting target recall within the latency budget; diminishing returns typically appear well before 20%.
Example 1: Input: "Our RAG system uses IVF with 128 clusters and centroid routing; recall is lower than we'd like at low latency budgets." Output: Replace centroid-based routing with a trained SupportNet that outputs k=128 scores jointly (one support-function estimate per cluster). Train on real user queries matched against ground-truth per-cluster max inner product values. At inference, route to top-s clusters by predicted support value instead of centroid distance. Expect disproportionate gains at low scan budgets — e.g., at 5% of the database scanned, learned routing recovers the true best-match partition ~90%+ of the time vs. ~87% for centroid routing; a small SupportNet can reach 95% routing recall at roughly half the latency of centroid routing.
Example 2: Input: "We have a production approximate search index we can't easily modify, but query latency to reach high recall is still too high." Output: Insert a KeyNet query rewriter in front of the existing index. Train KeyNet via regression on (query, ground-truth best key) pairs sampled from production query logs. At inference, pass the query through KeyNet first, then feed the predicted key vector (not the raw query) into the existing index as the search target. This moves the starting point much closer to the true answer, letting the unmodified index reach high recall with far less exploration — up to ~8x latency reduction to reach 99% recall vs. using the raw query.
- Sample training queries from the actual deployment query distribution, never from the key distribution — the value proposition depends on queries and keys not being interchangeable.
- Treat exhaustive ground-truth computation as an offline, amortized cost (once per training query, not per deployment query).
- Train separate specialized networks for support-value regression vs. best-key regression rather than relying on autodiff gradients of a single network — in practice this works better.
- Always report both recall-vs-latency (implementation-aware) and recall-vs-budget (implementation-agnostic) curves — they tell different stories.
- Use loosely-constrained ICNNs for SupportNet to respect convex structure; don't over-constrain KeyNet architecture — a consistency loss is sufficient.
- Sweep model size explicitly as a fraction of database size for a clean, reusable cost/accuracy tradeoff axis.
- Don't assume queries resemble keys geometrically — this assumption underlies naive centroid routing and is exactly what breaks down in RAG (short queries vs. long documents).
- Don't ignore NN inference overhead when reporting latency gains — always report net latency including model cost.
- Don't rely purely on centroid distance as a proxy for cluster quality — a cluster can hold an excellent match far from its own centroid, causing it to be skipped.
- Don't skip the budget-only (recall vs. fraction scanned) evaluation — latency numbers are hardware/implementation-dependent, but budget curves reveal the true algorithmic contribution.
- Don't default to the largest model size — benchmark the full XS→L sweep; mid-size models often capture most of the achievable gains.