Apple trains small neural nets to speed up RAG retrieval

Apple researchers have described training small neural networks, once and offline, to predict where a query's nearest match sits in a large vector database — a change to the retrieval step that the team says lowers search latency while improving accuracy. For anyone running retrieval-augmented generation at scale, that step is usually where cost and answer quality are decided.

What did Apple publish, and when?

According to Apple Machine Learning Research, which posted the work as a research highlight on November 5, 2026, the paper is titled "Amortizing Maximum Inner Product Search with Learned Support Functions" and was presented at ICML 2026. Its subject is maximum inner product search (MIPS), the primitive behind recommendation systems, search and RAG: given a query vector, the system looks through millions of candidate key vectors for the closest one, and the data point attached to that key is what gets used — in RAG, the best-suited document is prefixed to the question.

Apple's argument is about where to spend effort. Existing approximate methods build their indexing structures from the properties of the key set and only loosely assume that queries will resemble keys. In RAG that assumption breaks, since queries are short user-written questions while keys are long, curated passages. Deployed systems also serve a stable, well-characterized stream of queries, a signal that query-agnostic methods leave unused. The team therefore pays an upfront cost — training small networks for one fixed database — in exchange for cheaper, more accurate retrieval at inference time.

Why is vector search slow, and how does centroid routing miss results?

Cost climbs sharply once a database holds millions of keys with thousands of dimensions each, because comparing a query against every key is expensive. The standard shortcut splits the keys into k clusters, with k much smaller than the total n, keeps one representative centroid per cluster, compares the query only to those k centroids, then scans just the winning cluster. That leaves roughly n/k comparisons instead of n, a factor-of-k saving.

The catch is that centroid routing is not exact. A cluster can hold a key far from its own centroid yet close to the query; when a query lands near such an outlier, the cluster scores low and is skipped, so the true best match is never examined. Apple's first figure animates exactly that miss.

What is the support-function idea?

Apple notes that the mapping from a query to the largest inner product it achieves with any key in a set is a well-studied object in convex analysis: the support function of that set. Two properties do the work. It is convex, so it can be approximated by input convex neural networks (ICNNs), a network class designed for convex functions. And its gradient at a query returns the best-matching key itself.

That turns retrieval into regression. Sample queries, compute the true MIPS answer once by exhaustive search, store the pair, and train a network to reproduce the answer from the query. Apple's second figure plots the support function of each of four clusters as a query moves through embedding space — the quantity a router needs, rather than a single distance to a centroid.

What are SupportNet and KeyNet?

The paper describes two networks. SupportNet learns support functions directly with a loosely constrained ICNN, and Apple found it works best when one support function is computed per data cluster. Routing can then ask a sharper question than "how close is the centroid?" — namely "how good is the best key inside this cluster for this query?" A single SupportNet emits k scores at once and the largest routes the query, which in principle permits exact routing because each score accounts for every key in its cluster.

KeyNet instead outputs the best-matching key vector for a query. A consistency loss derived from Euler's theorem for homogeneous functions keeps its predictions geometrically grounded in the support function even though convexity is not enforced. Apple reports that training two specialized networks worked better than one network doing both jobs, and both are trained on samples from the query distribution that matters, with ground truth computed once by exhaustive search.

How much faster is retrieval, according to Apple?

The evaluation uses Natural Questions, a retrieval benchmark pairing general-knowledge questions with passages from a large document collection — roughly 2.6M unique passages in Apple's setup — partitioned into 128 clusters. The routers are 16-layer networks in four sizes, XS through L, corresponding to 1%, 5%, 10% and 20% of the database size. The metric is routing recall: how often the partition holding the exact best key is among those the router probes, with exhaustive search then run over the top-s clusters as s doubles from 1 to 128.

Two results stand out. Counting the cost of running the network, the small SupportNet reaches 95% routing recall about 1.9 times faster than centroid routing — 1.74 ms per query against 3.36 ms — and the large model reaches 97% in 1.74 ms. Ignoring network overhead and holding the scanned fraction of the database at 5%, learned routers find the correct partition about 90–93% of the time, up to 92.6% for the mid-size model, versus 87.5% for the centroid baseline. Every method reaches full recall once all 128 clusters are searched.

Does it work with an existing vector database?

In two different ways. SupportNet routing is presented as a replacement for the centroid-based routing inside existing inverted file (IVF) index systems. KeyNet is the less invasive option: the IVF index stays as it is, and the network rewrites the query first, predicting where the best match is likely to sit in embedding space so the index searches there instead of starting from the original query. KeyNet is trained by the same regression recipe. Apple evaluated KeyNets at the same four sizes, but the published text available to us breaks off partway through that recall–cost comparison, so those figures are not reported here.

What does this mean in practice, and how can you explore it in Grafy?

The clearest fit is retrieval over a fixed corpus with a predictable query mix — product search, or an assistant over a stable document set — where a one-off training run is amortized across a long stretch of traffic. The trade-off is that the networks are trained for one database and one query distribution, so a corpus refresh or a change in how people phrase questions implies retraining. The published evidence also covers a single public benchmark.

To read the paper itself, Grafy Reader can take the PDF and answer questions about the details that matter to your stack, such as how loosely the ICNN convexity constraint is enforced or what the Euler's-theorem consistency loss actually penalizes. To test the idea, Grafy Code can build a small IVF baseline plus a SupportNet-style router over your own embeddings and report recall against latency on your data — the same comparison Apple makes. For how running code in the studio works, see the AI coding studio walkthrough.

Frequently asked questions

Does SupportNet replace centroid routing in an existing IVF index?

According to Apple Machine Learning Research, yes. SupportNet routing is presented as a replacement for the centroid-based routing step in existing inverted file index systems, leaving the rest of the index in place. Rather than scoring clusters by distance to a centroid, the router outputs one score per cluster reflecting the best key any of them holds for that query, and the query is sent to the highest-scoring clusters for exhaustive search.

Can KeyNet be added in front of an index without rebuilding it?

Yes, and that is how Apple frames it. KeyNet works as a query rewriter in front of a standard approximate search index: it predicts where the best match is likely to be in embedding space, and the index is asked to look there rather than starting from the original query. Because the IVF index is already highly optimized, Apple says it finds the result faster with that head start.

Was the method tested on anything besides Natural Questions?

Not in the material Apple published. The reported experiments use Natural Questions, split into 128 clusters over roughly 2.6M unique passages, with 16-layer routers at four sizes. No other datasets are named, and the KeyNet recall–cost results are cut off mid-sentence in the available text, so the latency and recall figures above should be read as specific to that benchmark and setup.

---

Source: Hacker News — front page — Faster Retrieval for Better, More Efficient RAG