Apple researchers introduce amortized Maximum Inner Product Search (MIPS), a regression-based approach that trains small neural networks offline to accelerate retrieval in RAG and recommendation systems. MIPS finds the database key vector maximizing inner product with a query vector, but exhaustive search becomes prohibitive at millions of keys and thousands of dimensions. Standard inverted file index (IVF) methods use centroid-based routing, which ignores cluster spread and query-key distribution mismatch — queries are short user questions, keys are long curated passages. The paper "Amortizing Maximum Inner Product Search with Learned Support Functions" (ICML 2026) exploits the fact that the MIPS objective is the support function of the key set, a convex function whose gradient yields the optimal key. Two architectures are trained via supervised regression on queries sampled from the production distribution: SupportNet, an input convex neural network (ICNN) that learns the support function per cluster to enable exact routing, and KeyNet, which directly predicts the best-matching key vector using a consistency loss from Euler's theorem for homogeneous functions. Evaluated on Natural Questions (2.6M passages, 128 clusters) with FAISS IVF, SupportNet models (XS/L at 1%/20% DB size) achieve 95% routing recall 1.9× faster than centroid routing (1.74 ms vs 3.36 ms), and 92.6% recall scanning 5% of clusters vs 87.5% baseline. KeyNet-L reaches 99% recall in 0.23 ms vs 1.8 ms raw query (8× speedup), and 96.7% recall scanning 0.001% of DB vs 86% raw. Both integrate with existing IVF indices; SupportNet suits calibrated scoring, KeyNet excels at low-latency approximation. The approach amortizes one-time training cost for efficient, distribution-aware inference.
Read original
hackernews