Apple Researchers Propose Neural Network Method to Speed Up RAG Search

Researchers at Apple have detailed a new technique to optimize Maximum Inner Product Search (MIPS), a core component behind recommendation engines and retrieval-augmented generation (RAG) systems. As databases expand to millions of vectors with thousands of dimensions, running exhaustive comparisons becomes computationally expensive. Traditional approximate methods typically rely on clustering keys while ignoring the specific distribution of user queries, even though production systems usually serve a well-characterized user base.

Presented in a paper at ICML 2026 titled Amortizing Maximum Inner Product Search with Learned Support Functions, the method uses small neural networks trained offline for a fixed database of keys. Traditional MIPS splits keys into clusters and compares incoming queries against cluster centroids, a routing mechanism that frequently introduces errors. By treating the function that maps a query to its largest inner product as a support function from convex analysis, the researchers leverage input convex neural networks to approximate the function at a lower computational cost.

The approach utilizes a regression technique where neural networks are trained to mimic the ground-truth support function or its gradient by precomputing training examples of user queries paired with optimal MIPS outputs. This upfront amortization trades initial offline computation for faster and more accurate retrieval during inference time. What to watch next: adoption of amortized support functions in commercial search architectures.