Amortizing MIPS Computation Cost with Learned Support Functions

Maximum Inner Product Search (MIPS) is a core machine learning subroutine that finds the best match between a given query and keys within a database. While conventional methods must perform a search every time, this research proposes Amortized MIPS, an approach that trains a neural network to directly predict the solution in order to reduce the repeated MIPS cost for a fixed database.

The key insight is that the MIPS value function is the Support Function of the key set. Based on this, two complementary models are introduced.

• SupportNet: An input-convex neural network trained to regress the support function, acting as a cluster router that guides queries to the relevant database partition.
• KeyNet: A vector-valued network that directly regresses the optimal key, serving as a drop-in model that can be immediately substituted into existing indexing pipelines.

In experiments on the BEIR benchmark, SupportNet and KeyNet trained on document embeddings significantly improved IVF match rate in terms of computation (FLOPs), number of probes, and wall-clock time.