Glonce

Apple researchers train small networks to speed up RAG vector search

researchThe Glonce editorsHacker News🔥69Glonce rating69Impact55Novelty60Relevance95Surprise35Players50Credibility80

Maximum inner product search (MIPS) is the lookup at the core of recommendation, search and retrieval-augmented generation (RAG). Given a query vector, the system looks through millions of key vectors for the best match, and in RAG the document tied to that key is added in front of the user's question. As databases reach millions of keys and vectors reach thousands of dimensions, exhaustive comparison gets expensive. Standard approximate methods cluster or quantize the keys. Apple's researchers argue these methods loosely assume queries look like keys, or ignore what queries look like at all. That is a poor fit for RAG, where queries are short user questions and keys are long, curated passages. They add that production systems usually serve a stable, well-characterized query distribution that generic methods leave untapped.

In "Amortizing Maximum Inner Product Search with Learned Support Functions", presented at ICML 2026, researchers from Apple train small neural networks once, offline, for a fixed database of keys. The networks learn to route queries toward their best match faster. The trade is an upfront amortization cost in exchange for faster, more accurate retrieval at inference time.

The starting point is a piece of convex analysis. The function that maps a query to the largest inner product with any key in a set is the support function of that set. It is convex, so input convex neural networks (ICNNs) can approximate it, and its gradient at a query returns the best-matching key itself. The authors describe this structure as largely overlooked before. Learning it is plain regression: sample queries, compute the exact MIPS answer for each (either the best score or the best key) by exhaustive search once, store it, and train a network to reproduce it. Training examples come from the query distribution that matters.

The authors found that two specialized networks work better than one network doing both jobs. SupportNet learns the support function directly with a loosely constrained ICNN, and it shines when one support function is learned per data cluster to improve routing. KeyNet outputs the best-matching key vector directly from a query. A consistency loss derived from Euler's theorem for homogeneous functions keeps its predictions grounded in the support function, though no convexity is enforced.

SupportNet replaces the centroid routing used in existing inverted file index (IVF) systems. Centroid routing compares a query with one representative point per cluster, which can miss a cluster holding an outlier key that sits far from its centroid but close to the query. A SupportNet outputs k scores at once, one per cluster, and answers the question: "How good is the best key inside this cluster for this query?" The largest scores decide where to search. KeyNet works differently: it is a query rewriter placed in front of a standard approximate index, predicting where the best match likely sits in embedding space so the unchanged IVF system searches there instead of at the original query.

The experiments use Natural Questions, a retrieval benchmark with roughly 2.6M unique passages in the authors' setup, partitioned into 128 clusters and served through a FAISS IVF index. Models come in four sizes, XS, S, M and L, at 1%, 5%, 10% and 20% of the database size. Routing recall is the share of queries whose true partition, the one holding the exact best key, is among those probed.

For SupportNet, latency figures include the cost of running the networks. The small model reaches 95% routing recall about 1.9× faster than centroid routing (1.74 ms versus 3.36 ms), and the large model hits 97% in 1.74 ms. A second view ignores network overhead and looks only at routing quality. Scanning 5% of the database, the learned routers find the true partition about 90 to 93% of the time, up to 92.6% for the mid-size M model, against 87.5% for centroid routing.

For KeyNet, reaching 99% recall takes up to about 8× less wall-clock time with the largest model (0.23 ms versus 1.8 ms for the raw query). In the quality-only view, scanning 0.001% of the database retrieves the true key 89 to 97% of the time, up to 96.7% for the largest model, against 86% for the raw query. Reaching 99% recall takes up to about 11× fewer keys scanned (0.009% of the database versus 0.10% for the raw query). The authors say Natural Questions is one of several datasets they studied and the trend is consistent across them. They also show five example questions with the top 3 hits for the raw query versus the KeyNet-mapped query.

Their practical advice: use KeyNet when inference speed is the bottleneck and you only need to approximate the best key. Use SupportNet when you also want a calibrated score for re-ranking, thresholding or downstream use, or when compute is tight at small model sizes, where it stays competitive. Both train similarly, plug into existing approximate-search infrastructure, and are reasonably robust to mild distribution shift in query patterns. The article text available here is cut off mid-sentence in its closing remarks on very large-scale systems.

Key facts

  • Apple researchers' ICML 2026 paper trains small neural networks once, offline, for a fixed key database, so that maximum inner product search (the lookup step in RAG) routes queries faster.
  • SupportNet replaces centroid routing in IVF indexes; on Natural Questions (roughly 2.6M passages, 128 clusters) the small model reaches 95% routing recall about 1.9× faster than centroids (1.74 ms vs 3.36 ms).
  • KeyNet rewrites the query in front of a standard index: reaching 99% recall takes up to about 8× less time with the largest model (0.23 ms vs 1.8 ms for the raw query) and up to about 11× fewer keys scanned.
  • Quality-only views ignore network overhead: scanning 5% of the database, SupportNet recovers the true partition up to 92.6% of the time vs 87.5% for centroids; scanning 0.001%, the largest KeyNet reaches 96.7% vs 86% for the raw query.
  • The authors say the trend holds across several datasets; the visible text reports retrieval recall and latency, not end-to-end RAG answer quality.

Why it matters

Every RAG system depends on a fast, accurate nearest-match lookup, and the cost of that lookup grows with the database. The paper's argument is that generic indexes ignore a useful signal: production systems often serve a stable, well-characterized stream of queries, and RAG queries (short questions) look very different from the keys (long, curated passages). Training small networks on that query distribution is a way to spend an upfront, one-off cost for cheaper searches later. The authors also point to a mathematical link they say was largely overlooked: MIPS is the support function of the key set, a convex object that neural networks can approximate.

Who it affects

Engineers who run vector search or RAG over a large, fixed corpus with a predictable stream of user questions, since the method is built for a database that does not change and a query pattern that can be sampled. It also bears on teams that already run IVF indexes such as FAISS, because both models are described as integrating with existing approximate-search infrastructure.

How to use it

The authors give a rule of thumb. Pick KeyNet when inference speed is the bottleneck and approximating the best key is enough; it sits in front of your existing index as a query rewriter and leaves the index untouched. Pick SupportNet when you want a calibrated score for re-ranking or thresholding, or when compute is tight at small model sizes; it replaces centroid routing in an IVF system. Both are trained by sampling queries from your expected distribution, computing the exact best match once by exhaustive search, and fitting a network by regression. Model sizes in the paper range from 1% to 20% of the database size.

How solid is it

The paper was presented at ICML 2026, and the article is the authors' own summary on Apple's machine learning site. The headline numbers come from one dataset, Natural Questions with roughly 2.6M passages, served through a FAISS IVF index; the authors say the trend is consistent across other datasets they studied but the visible text gives no figures for them. Some figures include the cost of running the networks (the latency results: 1.9× for SupportNet, about 8× for KeyNet), while others do not (the 92.6% vs 87.5% routing recall, the 96.7% vs 86% KeyNet recall, and the 11× key-count reduction). The article text available here is cut off before its closing remarks.

Risks and caveats

The approach pays an upfront training cost, and the visible text does not quantify it. It is tuned to a fixed key database and a known query distribution, and the authors say the models are reasonably robust to mild shift in query patterns, which is not the same as robust to large shifts. The reported latencies come without stated hardware, software stack or batch size. The results measure retrieval recall and latency, not end-to-end RAG answer quality. The text makes no claim that the methods are deployed in any Apple product.

“How good is the best key inside this cluster for this query?”

— From Apple's article on the paper, describing what a SupportNet can answer