Amortizing Maximum Inner Product Search with Learned Support Functions

TL;DR

This paper introduces a learning-based method for maximum inner product search (MIPS) that uses neural networks to directly predict solutions, significantly reducing computational costs by leveraging the support function of the key set.

Contribution

It proposes novel neural network models, SupportNet and KeyNet, to approximate the support function and optimal keys, enabling efficient amortized MIPS with theoretical guarantees.

Findings

SupportNet and KeyNet achieve high match rates in experiments.

The methods enable efficient database compression tailored to specific query distributions.

The approach links neural network approximation to convex analysis and homogeneous functions.

Abstract

Maximum inner product search (MIPS) is a crucial subroutine in machine learning, requiring the identification of key vectors that align best with a given query. We propose amortized MIPS: a learning-based approach that trains neural networks to directly predict MIPS solutions, amortizing the computational cost of matching queries (drawn from a fixed distribution) to a fixed set of keys. Our key insight is that the MIPS value function, the maximal inner product between a query and keys, is also known as the support function of the set of keys. Support functions are convex, 1-homogeneous and their gradient w.r.t. the query is exactly the optimal key in the database. We approximate the support function using two complementary approaches: (1) we train an input-convex neural network (SupportNet) to model the support function directly; the optimal key can be recovered via (autodiff) gradient computation, and (2) we regress directly the optimal key from the query using a vector valued network (KeyNet), bypassing gradient computation entirely at inference time. To learn a SupportNet, we combine score regression with gradient matching losses, and propose homogenization wrappers that enforce the positive 1-homogeneity of a neural network, theoretically linking function values to gradients. To train a KeyNet, we introduce a score consistency loss derived from the Euler theorem for homogeneous functions. Our experiments show that learned SupportNet or KeyNet achieve high match rates and open up new directions to compress databases with a specific query distribution in mind.