Amortizing Maximum Inner Product Search with Learned Support Functions

Maximum inner product search (MIPS) is a key subroutine in machine learning, but repeatedly solving it for queries from a known distribution over a fixed database is computationally expensive.

The authors propose amortized MIPS, a regression-based approach that trains neural networks to directly predict MIPS solutions, reducing the cost of repeated search.

The key insight is that the MIPS value function is the support function of the set of keys, a convex function whose gradient yields the optimal key.

This motivates two models: SupportNet, an input-convex neural network that regresses the support function and can serve as a cluster router, and KeyNet, a vector-valued network that directly regresses the optimal key and can be used as a drop-in replacement for queries in off-the-shelf indexing pipelines.

Experiments on the BEIR benchmark show that for document embeddings, learned SupportNets and KeyNets significantly improve IVF match rates when accounting for compute effort, measured in FLOPs, number of probes, or wall-clock time.

This work demonstrates that leveraging convex analysis and amortized inference can make large-scale similarity search more efficient without sacrificing accuracy.