Apple Researchers Train Neural Networks to Speed Up RAG Retrieval

According to Apple Machine Learning:

In a paper presented at ICML 2026, Apple researchers introduced a technique called amortized maximum inner product search (MIPS) to improve retrieval-augmented generation (RAG) and search systems. Traditional systems often rely on geometric approximations, such as clustering keys and comparing queries to cluster centroids, which can lead to errors because queries (short user questions) and keys (long curated texts) differ significantly. To solve this, the researchers train small neural networks offline—specifically SupportNet and KeyNet—using supervised learning on typical query distributions.

SupportNet learns the mathematical support function of data clusters to route queries more accurately, helping the system find the right partition up to 1.9 times faster than traditional centroid routing. Meanwhile, KeyNet acts as a query rewriter that transforms incoming questions directly into predicted key vectors, allowing existing index systems to retrieve true keys up to 11 times fewer database scans. This method effectively trades a one-time training cost for significantly faster inference, demonstrating how knowing the query distribution in advance can dramatically optimize computational resources in production machine learning systems.